Browse Source

Add tool message types; show tool specific information in webview; separate command from output; add abort button to task card

Saoud Rizwan 1 year ago
parent
commit
d63aef015a

+ 40 - 32
src/ClaudeDev.ts

@@ -9,7 +9,7 @@ import * as path from "path"
 import { serializeError } from "serialize-error"
 import { serializeError } from "serialize-error"
 import { DEFAULT_MAX_REQUESTS_PER_TASK } from "./shared/Constants"
 import { DEFAULT_MAX_REQUESTS_PER_TASK } from "./shared/Constants"
 import { Tool, ToolName } from "./shared/Tool"
 import { Tool, ToolName } from "./shared/Tool"
-import { ClaudeAsk, ClaudeSay, ExtensionMessage } from "./shared/ExtensionMessage"
+import { ClaudeAsk, ClaudeSay, ClaudeSayTool, ExtensionMessage } from "./shared/ExtensionMessage"
 import * as vscode from "vscode"
 import * as vscode from "vscode"
 import pWaitFor from "p-wait-for"
 import pWaitFor from "p-wait-for"
 import { ClaudeAskResponse } from "./shared/WebviewMessage"
 import { ClaudeAskResponse } from "./shared/WebviewMessage"
@@ -214,7 +214,7 @@ export class ClaudeDev {
 		await this.providerRef.deref()?.postStateToWebview()
 		await this.providerRef.deref()?.postStateToWebview()
 
 
 		// Get all relevant context for the task
 		// Get all relevant context for the task
-		const filesInCurrentDir = await this.listFiles(".")
+		const filesInCurrentDir = await this.listFiles(".", false)
 
 
 		// This first message kicks off a task, it is not included in every subsequent message. This is a good place to give all the relevant context to a task, instead of having Claude request for it using tools.
 		// This first message kicks off a task, it is not included in every subsequent message. This is a good place to give all the relevant context to a task, instead of having Claude request for it using tools.
 		let userPrompt = `# Task
 		let userPrompt = `# Task
@@ -235,7 +235,7 @@ ${filesInCurrentDir}`
 ${activeEditorContents}`
 ${activeEditorContents}`
 		}
 		}
 
 
-		await this.say("text", userPrompt)
+		await this.say("text", task)
 
 
 		let totalInputTokens = 0
 		let totalInputTokens = 0
 		let totalOutputTokens = 0
 		let totalOutputTokens = 0
@@ -255,10 +255,10 @@ ${activeEditorContents}`
 				//this.say("task_completed", `Task completed. Total API usage cost: ${totalCost}`)
 				//this.say("task_completed", `Task completed. Total API usage cost: ${totalCost}`)
 				break
 				break
 			} else {
 			} else {
-				this.say(
-					"tool",
-					"Claude responded with only text blocks but has not called attempt_completion yet. Forcing him to continue with task..."
-				)
+				// this.say(
+				// 	"tool",
+				// 	"Claude responded with only text blocks but has not called attempt_completion yet. Forcing him to continue with task..."
+				// )
 				userPrompt =
 				userPrompt =
 					"Ask yourself if you have completed the user's task. If you have, use the attempt_completion tool, otherwise proceed to the next step. (This is an automated message, so do not respond to it conversationally. Just proceed with the task.)"
 					"Ask yourself if you have completed the user's task. If you have, use the attempt_completion tool, otherwise proceed to the next step. (This is an automated message, so do not respond to it conversationally. Just proceed with the task.)"
 			}
 			}
@@ -305,38 +305,46 @@ ${activeEditorContents}`
 				const diffResult = diff.createPatch(filePath, originalContent, newContent)
 				const diffResult = diff.createPatch(filePath, originalContent, newContent)
 				if (diffResult) {
 				if (diffResult) {
 					await fs.writeFile(filePath, newContent)
 					await fs.writeFile(filePath, newContent)
+					this.say("tool", JSON.stringify({ tool: "editedExistingFile", path: filePath, diff: diffResult } as ClaudeSayTool))
 					return `Changes applied to ${filePath}:\n${diffResult}`
 					return `Changes applied to ${filePath}:\n${diffResult}`
 				} else {
 				} else {
+					this.say("tool", JSON.stringify({ tool: "editedExistingFile", path: filePath } as ClaudeSayTool))
 					return `Tool succeeded, however there were no changes detected to ${filePath}`
 					return `Tool succeeded, however there were no changes detected to ${filePath}`
 				}
 				}
 			} else {
 			} else {
 				await fs.mkdir(path.dirname(filePath), { recursive: true })
 				await fs.mkdir(path.dirname(filePath), { recursive: true })
 				await fs.writeFile(filePath, newContent)
 				await fs.writeFile(filePath, newContent)
+				this.say("tool", JSON.stringify({ tool: "newFileCreated", path: filePath, content: newContent } as ClaudeSayTool))
 				return `New file created and content written to ${filePath}`
 				return `New file created and content written to ${filePath}`
 			}
 			}
 		} catch (error) {
 		} catch (error) {
 			const errorString = `Error writing file: ${JSON.stringify(serializeError(error))}`
 			const errorString = `Error writing file: ${JSON.stringify(serializeError(error))}`
-			this.say("error", errorString)
+			this.say("error", JSON.stringify(serializeError(error)))
 			return errorString
 			return errorString
 		}
 		}
 	}
 	}
 
 
 	async readFile(filePath: string): Promise<string> {
 	async readFile(filePath: string): Promise<string> {
 		try {
 		try {
-			return await fs.readFile(filePath, "utf-8")
+			const content = await fs.readFile(filePath, "utf-8")
+			this.say("tool", JSON.stringify({ tool: "readFile", path: filePath } as ClaudeSayTool))
+			return content
 		} catch (error) {
 		} catch (error) {
 			const errorString = `Error reading file: ${JSON.stringify(serializeError(error))}`
 			const errorString = `Error reading file: ${JSON.stringify(serializeError(error))}`
-			this.say("error", errorString)
+			this.say("error", JSON.stringify(serializeError(error)))
 			return errorString
 			return errorString
 		}
 		}
 	}
 	}
 
 
-	async listFiles(dirPath: string): Promise<string> {
+	async listFiles(dirPath: string, shouldLog: boolean = true): Promise<string> {
 		// If the extension is run without a workspace open, we are in the root directory and don't want to list all files since it would prompt for permission to access everything
 		// If the extension is run without a workspace open, we are in the root directory and don't want to list all files since it would prompt for permission to access everything
 		const cwd = process.cwd()
 		const cwd = process.cwd()
 		const root = process.platform === "win32" ? path.parse(cwd).root : "/"
 		const root = process.platform === "win32" ? path.parse(cwd).root : "/"
 		const isRoot = cwd === root
 		const isRoot = cwd === root
 		if (isRoot) {
 		if (isRoot) {
+			if (shouldLog) {
+				this.say("tool", JSON.stringify({ tool: "listFiles", path: dirPath } as ClaudeSayTool))
+			}
 			return "Currently in the root directory. Cannot list all files."
 			return "Currently in the root directory. Cannot list all files."
 		}
 		}
 
 
@@ -348,19 +356,19 @@ ${activeEditorContents}`
 			}
 			}
 			// * globs all files in one dir, ** globs files in nested directories
 			// * globs all files in one dir, ** globs files in nested directories
 			const entries = await glob("*", options)
 			const entries = await glob("*", options)
+			if (shouldLog) {
+				this.say("tool", JSON.stringify({ tool: "listFiles", path: dirPath } as ClaudeSayTool))
+			}
 			return entries.slice(0, 500).join("\n") // truncate to 500 entries
 			return entries.slice(0, 500).join("\n") // truncate to 500 entries
 		} catch (error) {
 		} catch (error) {
 			const errorString = `Error listing files and directories: ${JSON.stringify(serializeError(error))}`
 			const errorString = `Error listing files and directories: ${JSON.stringify(serializeError(error))}`
-			this.say("error", errorString)
+			this.say("error", JSON.stringify(serializeError(error)))
 			return errorString
 			return errorString
 		}
 		}
 	}
 	}
 
 
 	async executeCommand(command: string): Promise<string> {
 	async executeCommand(command: string): Promise<string> {
-		const { response } = await this.ask(
-			"command",
-			`Claude wants to execute the following command:\n${command}\nDo you approve?`
-		)
+		const { response } = await this.ask("command", command)
 		if (response !== "yesButtonTapped") {
 		if (response !== "yesButtonTapped") {
 			return "Command execution was not approved by the user."
 			return "Command execution was not approved by the user."
 		}
 		}
@@ -378,7 +386,7 @@ ${activeEditorContents}`
 			const error = e as any
 			const error = e as any
 			let errorMessage = error.message || JSON.stringify(serializeError(error))
 			let errorMessage = error.message || JSON.stringify(serializeError(error))
 			const errorString = `Error executing command:\n${errorMessage}`
 			const errorString = `Error executing command:\n${errorMessage}`
-			this.say("error", errorString)
+			this.say("error", errorMessage)
 			return errorString
 			return errorString
 		}
 		}
 	}
 	}
@@ -414,7 +422,7 @@ ${activeEditorContents}`
 		if (this.requestCount >= this.maxRequestsPerTask) {
 		if (this.requestCount >= this.maxRequestsPerTask) {
 			const { response } = await this.ask(
 			const { response } = await this.ask(
 				"request_limit_reached",
 				"request_limit_reached",
-				`\nClaude has exceeded ${this.maxRequestsPerTask} requests for this task! Would you like to reset the count and proceed?:`
+				`Claude Dev has reached the maximum number of requests for this task. Would you like to reset the count and allow him to proceed?`
 			)
 			)
 
 
 			if (response === "yesButtonTapped") {
 			if (response === "yesButtonTapped") {
@@ -434,7 +442,7 @@ ${activeEditorContents}`
 		}
 		}
 
 
 		try {
 		try {
-			await this.say("api_req_started", JSON.stringify(userContent))
+			await this.say("api_req_started", JSON.stringify({ request: userContent }))
 			const response = await this.client.messages.create({
 			const response = await this.client.messages.create({
 				model: "claude-3-5-sonnet-20240620", // https://docs.anthropic.com/en/docs/about-claude/models
 				model: "claude-3-5-sonnet-20240620", // https://docs.anthropic.com/en/docs/about-claude/models
 				max_tokens: 4096,
 				max_tokens: 4096,
@@ -448,7 +456,7 @@ ${activeEditorContents}`
 			let assistantResponses: Anthropic.Messages.ContentBlock[] = []
 			let assistantResponses: Anthropic.Messages.ContentBlock[] = []
 			let inputTokens = response.usage.input_tokens
 			let inputTokens = response.usage.input_tokens
 			let outputTokens = response.usage.output_tokens
 			let outputTokens = response.usage.output_tokens
-			await this.say("api_req_finished", this.calculateApiCost(inputTokens, outputTokens).toString())
+			await this.say("api_req_finished", JSON.stringify({ tokensIn: inputTokens, tokensOut: outputTokens, cost: this.calculateApiCost(inputTokens, outputTokens) }))
 
 
 			// A response always returns text content blocks (it's just that before we were iterating over the completion_attempt response before we could append text response, resulting in bug)
 			// A response always returns text content blocks (it's just that before we were iterating over the completion_attempt response before we could append text response, resulting in bug)
 			for (const contentBlock of response.content) {
 			for (const contentBlock of response.content) {
@@ -470,10 +478,10 @@ ${activeEditorContents}`
 						attemptCompletionBlock = contentBlock
 						attemptCompletionBlock = contentBlock
 					} else {
 					} else {
 						const result = await this.executeTool(toolName, toolInput)
 						const result = await this.executeTool(toolName, toolInput)
-						this.say(
-							"tool",
-							`\nTool Used: ${toolName}\nTool Input: ${JSON.stringify(toolInput)}\nTool Result: ${result}`
-						)
+						// this.say(
+						// 	"tool",
+						// 	`\nTool Used: ${toolName}\nTool Input: ${JSON.stringify(toolInput)}\nTool Result: ${result}`
+						// )
 						toolResults.push({ type: "tool_result", tool_use_id: toolUseId, content: result })
 						toolResults.push({ type: "tool_result", tool_use_id: toolUseId, content: result })
 					}
 					}
 				}
 				}
@@ -483,7 +491,7 @@ ${activeEditorContents}`
 				this.conversationHistory.push({ role: "assistant", content: assistantResponses })
 				this.conversationHistory.push({ role: "assistant", content: assistantResponses })
 			} else {
 			} else {
 				// this should never happen! it there's no assistant_responses, that means we got no text or tool_use content blocks from API which we should assume is an error
 				// this should never happen! it there's no assistant_responses, that means we got no text or tool_use content blocks from API which we should assume is an error
-				this.say("error", "Error: No assistant responses found in API response!")
+				this.say("error", "Unexpected Error: No assistant messages were found in the API response")
 				this.conversationHistory.push({
 				this.conversationHistory.push({
 					role: "assistant",
 					role: "assistant",
 					content: [{ type: "text", text: "Failure: I did not have a response to provide." }],
 					content: [{ type: "text", text: "Failure: I did not have a response to provide." }],
@@ -499,12 +507,12 @@ ${activeEditorContents}`
 					attemptCompletionBlock.name as ToolName,
 					attemptCompletionBlock.name as ToolName,
 					attemptCompletionBlock.input
 					attemptCompletionBlock.input
 				)
 				)
-				this.say(
-					"tool",
-					`\nattempt_completion Tool Used: ${attemptCompletionBlock.name}\nTool Input: ${JSON.stringify(
-						attemptCompletionBlock.input
-					)}\nTool Result: ${result}`
-				)
+				// this.say(
+				// 	"tool",
+				// 	`\nattempt_completion Tool Used: ${attemptCompletionBlock.name}\nTool Input: ${JSON.stringify(
+				// 		attemptCompletionBlock.input
+				// 	)}\nTool Result: ${result}`
+				// )
 				if (result === "") {
 				if (result === "") {
 					didCompleteTask = true
 					didCompleteTask = true
 					result = "The user is satisfied with the result."
 					result = "The user is satisfied with the result."
@@ -539,7 +547,7 @@ ${activeEditorContents}`
 			return { didCompleteTask, inputTokens, outputTokens }
 			return { didCompleteTask, inputTokens, outputTokens }
 		} catch (error) {
 		} catch (error) {
 			// only called if the API request fails (executeTool errors are returned back to claude)
 			// only called if the API request fails (executeTool errors are returned back to claude)
-			this.say("error", `Error calling Claude API: ${JSON.stringify(serializeError(error))}`)
+			this.say("error", JSON.stringify(serializeError(error)))
 			return { didCompleteTask: true, inputTokens: 0, outputTokens: 0 }
 			return { didCompleteTask: true, inputTokens: 0, outputTokens: 0 }
 		}
 		}
 	}
 	}

+ 8 - 1
src/shared/ExtensionMessage.ts

@@ -17,4 +17,11 @@ export interface ClaudeMessage {
 }
 }
 
 
 export type ClaudeAsk = "request_limit_reached" | "followup" | "command" | "completion_result"
 export type ClaudeAsk = "request_limit_reached" | "followup" | "command" | "completion_result"
-export type ClaudeSay = "task" | "error" | "api_req_started" | "api_req_finished" | "text" | "tool" | "command_output" | "completion_result"
+export type ClaudeSay = "task" | "error" | "api_req_started" | "api_req_finished" | "text" | "tool" | "command_output" | "completion_result"
+
+export interface ClaudeSayTool {
+    tool: "editedExistingFile" | "newFileCreated" | "readFile" | "listFiles"
+    path?: string
+    diff?: string
+    content?: string
+}

+ 1 - 1
src/shared/WebviewMessage.ts

@@ -1,5 +1,5 @@
 export interface WebviewMessage {
 export interface WebviewMessage {
-    type: "apiKey" | "maxRequestsPerTask" | "webviewDidLaunch" | "newTask" | "askResponse"
+    type: "apiKey" | "maxRequestsPerTask" | "webviewDidLaunch" | "newTask" | "askResponse" | "abortTask"
     text?: string
     text?: string
     askResponse?: ClaudeAskResponse
     askResponse?: ClaudeAskResponse
 }
 }

+ 79 - 21
webview-ui/src/components/ChatRow.tsx

@@ -1,6 +1,7 @@
 import React, { useState } from "react"
 import React, { useState } from "react"
-import { ClaudeMessage, ClaudeAsk, ClaudeSay } from "@shared/ExtensionMessage"
+import { ClaudeMessage, ClaudeAsk, ClaudeSay, ClaudeSayTool } from "@shared/ExtensionMessage"
 import { VSCodeButton, VSCodeProgressRing, VSCodeBadge } from "@vscode/webview-ui-toolkit/react"
 import { VSCodeButton, VSCodeProgressRing, VSCodeBadge } from "@vscode/webview-ui-toolkit/react"
+import { COMMAND_OUTPUT_STRING } from "../utilities/combineCommandSequences"
 
 
 interface ChatRowProps {
 interface ChatRowProps {
 	message: ClaudeMessage
 	message: ClaudeMessage
@@ -92,6 +93,7 @@ const ChatRow: React.FC<ChatRowProps> = ({ message }) => {
 
 
 		const contentStyle: React.CSSProperties = {
 		const contentStyle: React.CSSProperties = {
 			margin: 0,
 			margin: 0,
+			whiteSpace: "pre-line",
 		}
 		}
 
 
 		switch (message.type) {
 		switch (message.type) {
@@ -116,18 +118,58 @@ const ChatRow: React.FC<ChatRowProps> = ({ message }) => {
 					case "api_req_finished":
 					case "api_req_finished":
 						return null // Hide this message type
 						return null // Hide this message type
 					case "tool":
 					case "tool":
+						//const tool = JSON.parse(message.text || "{}") as ClaudeSayTool
+						const tool: ClaudeSayTool = {
+							tool: "editedExistingFile",
+							path: "/path/to/file",
+						}
+						switch (tool.tool) {
+							case "editedExistingFile":
+								return (
+									<>
+										<div style={headerStyle}>
+											{icon}
+											Edited File
+										</div>
+										<p>Path: {tool.path!}</p>
+										<p>{tool.diff!}</p>
+									</>
+								)
+							case "newFileCreated":
+								return (
+									<>
+										<div style={headerStyle}>
+											{icon}
+											Created New File
+										</div>
+										<p>Path: {tool.path!}</p>
+										<p>{tool.content!}</p>
+									</>
+								)
+							case "readFile":
+								return (
+									<>
+										<div style={headerStyle}>
+											{icon}
+											Read File
+										</div>
+										<p>Path: {tool.path!}</p>
+									</>
+								)
+							case "listFiles":
+								return (
+									<>
+										<div style={headerStyle}>
+											{icon}
+											Viewed Directory
+										</div>
+										<p>Path: {tool.path!}</p>
+									</>
+								)
+						}
+						break
 					case "text":
 					case "text":
-						return (
-							<>
-								{title && (
-									<div style={headerStyle}>
-										{icon}
-										{title}
-									</div>
-								)}
-								<p style={contentStyle}>{message.text}</p>
-							</>
-						)
+						return <p style={contentStyle}>{message.text}</p>
 					case "error":
 					case "error":
 						return (
 						return (
 							<>
 							<>
@@ -167,6 +209,7 @@ const ChatRow: React.FC<ChatRowProps> = ({ message }) => {
 							</>
 							</>
 						)
 						)
 				}
 				}
+				break
 			case "ask":
 			case "ask":
 				switch (message.ask) {
 				switch (message.ask) {
 					case "request_limit_reached":
 					case "request_limit_reached":
@@ -177,12 +220,23 @@ const ChatRow: React.FC<ChatRowProps> = ({ message }) => {
 									{title}
 									{title}
 								</div>
 								</div>
 								<p style={{ ...contentStyle, color: "var(--vscode-errorForeground)" }}>
 								<p style={{ ...contentStyle, color: "var(--vscode-errorForeground)" }}>
-									Your task has reached the maximum request limit (maxRequestsPerTask, you can change
-									this in settings). Do you want to keep going or start a new task?
+									{message.text}
 								</p>
 								</p>
 							</>
 							</>
 						)
 						)
 					case "command":
 					case "command":
+						const splitMessage = (text: string) => {
+							const outputIndex = text.indexOf(COMMAND_OUTPUT_STRING)
+							if (outputIndex === -1) {
+								return { command: text, output: "" }
+							}
+							return {
+								command: text.slice(0, outputIndex).trim(),
+								output: text.slice(outputIndex + COMMAND_OUTPUT_STRING.length).trim(),
+							}
+						}
+
+						const { command, output } = splitMessage(message.text || "")
 						return (
 						return (
 							<>
 							<>
 								<div style={headerStyle}>
 								<div style={headerStyle}>
@@ -190,10 +244,16 @@ const ChatRow: React.FC<ChatRowProps> = ({ message }) => {
 									{title}
 									{title}
 								</div>
 								</div>
 								<div style={contentStyle}>
 								<div style={contentStyle}>
-									<p>Claude would like to run this command. Do you allow this?</p>
-									<pre style={contentStyle}>
-										<code>{message.text}</code>
-									</pre>
+									<p style={contentStyle}>Claude Dev wants to execute the following command:</p>
+									<p style={contentStyle}>{command}</p>
+									{output && (
+										<>
+											<p style={{ ...contentStyle, fontWeight: "bold" }}>
+												{COMMAND_OUTPUT_STRING}
+											</p>
+											<p style={contentStyle}>{output}</p>
+										</>
+									)}
 								</div>
 								</div>
 							</>
 							</>
 						)
 						)
@@ -240,9 +300,7 @@ const ChatRow: React.FC<ChatRowProps> = ({ message }) => {
 			}}>
 			}}>
 			{renderContent()}
 			{renderContent()}
 			{isExpanded && message.say === "api_req_started" && (
 			{isExpanded && message.say === "api_req_started" && (
-				<pre style={{ marginTop: "10px" }}>
-					<code>{message.text}</code>
-				</pre>
+				<p style={{ marginTop: "10px" }}>{JSON.stringify(JSON.parse(message.text || "{}").request)}</p>
 			)}
 			)}
 		</div>
 		</div>
 	)
 	)

+ 21 - 1
webview-ui/src/components/ChatView.tsx

@@ -33,7 +33,11 @@ const ChatView = ({ messages }: ChatViewProps) => {
 
 
 	const scrollToBottom = (instant: boolean = false) => {
 	const scrollToBottom = (instant: boolean = false) => {
 		// https://stackoverflow.com/questions/11039885/scrollintoview-causing-the-whole-page-to-move
 		// https://stackoverflow.com/questions/11039885/scrollintoview-causing-the-whole-page-to-move
-		(messagesEndRef.current as any)?.scrollIntoView({ behavior: instant ? "instant" : "smooth", block: "nearest", inline: "start" })
+		;(messagesEndRef.current as any)?.scrollIntoView({
+			behavior: instant ? "instant" : "smooth",
+			block: "nearest",
+			inline: "start",
+		})
 	}
 	}
 
 
 	const handlePrimaryButtonClick = () => {
 	const handlePrimaryButtonClick = () => {
@@ -49,8 +53,19 @@ const ChatView = ({ messages }: ChatViewProps) => {
 		setSecondaryButtonText(undefined)
 		setSecondaryButtonText(undefined)
 	}
 	}
 
 
+	// scroll to bottom when new message is added
+	const visibleMessages = useMemo(
+		() =>
+			modifiedMessages.filter(
+				(message) => !(message.type === "ask" && message.ask === "completion_result" && message.text === "")
+			),
+		[modifiedMessages]
+	)
 	useEffect(() => {
 	useEffect(() => {
 		scrollToBottom()
 		scrollToBottom()
+	}, [visibleMessages.length])
+
+	useEffect(() => {
 		// if last message is an ask, show user ask UI
 		// if last message is an ask, show user ask UI
 
 
 		// if user finished a task, then start a new task with a new conversation history since in this moment that the extension is waiting for user response, the user could close the extension and the conversation history would be lost.
 		// if user finished a task, then start a new task with a new conversation history since in this moment that the extension is waiting for user response, the user could close the extension and the conversation history would be lost.
@@ -110,6 +125,10 @@ const ChatView = ({ messages }: ChatViewProps) => {
 		}
 		}
 	}
 	}
 
 
+	const handleTaskCloseButtonClick = () => {
+		vscode.postMessage({ type: "abortTask" })
+	}
+
 	useEffect(() => {
 	useEffect(() => {
 		if (textAreaRef.current && !textAreaHeight) {
 		if (textAreaRef.current && !textAreaHeight) {
 			setTextAreaHeight(textAreaRef.current.offsetHeight)
 			setTextAreaHeight(textAreaRef.current.offsetHeight)
@@ -158,6 +177,7 @@ const ChatView = ({ messages }: ChatViewProps) => {
 				tokensIn={apiMetrics.totalTokensIn}
 				tokensIn={apiMetrics.totalTokensIn}
 				tokensOut={apiMetrics.totalTokensOut}
 				tokensOut={apiMetrics.totalTokensOut}
 				totalCost={apiMetrics.totalCost}
 				totalCost={apiMetrics.totalCost}
+				onClose={handleTaskCloseButtonClick}
 			/>
 			/>
 			<div
 			<div
 				className="scrollable"
 				className="scrollable"

+ 35 - 18
webview-ui/src/components/TaskHeader.tsx

@@ -1,32 +1,45 @@
 import React, { useState } from "react"
 import React, { useState } from "react"
 import TextTruncate from "react-text-truncate"
 import TextTruncate from "react-text-truncate"
+import { VSCodeButton } from "@vscode/webview-ui-toolkit/react"
 
 
 interface TaskHeaderProps {
 interface TaskHeaderProps {
 	taskText: string
 	taskText: string
 	tokensIn: number
 	tokensIn: number
 	tokensOut: number
 	tokensOut: number
 	totalCost: number
 	totalCost: number
+	onClose: () => void
 }
 }
 
 
-const TaskHeader: React.FC<TaskHeaderProps> = ({ taskText, tokensIn, tokensOut, totalCost }) => {
+const TaskHeader: React.FC<TaskHeaderProps> = ({ taskText, tokensIn, tokensOut, totalCost, onClose }) => {
 	const [isExpanded, setIsExpanded] = useState(false)
 	const [isExpanded, setIsExpanded] = useState(false)
 	const toggleExpand = () => setIsExpanded(!isExpanded)
 	const toggleExpand = () => setIsExpanded(!isExpanded)
 
 
 	return (
 	return (
-		<div
-			style={{
-				padding: "15px 15px 10px 15px",
-			}}>
+		<div style={{ padding: "15px 15px 10px 15px" }}>
 			<div
 			<div
 				style={{
 				style={{
 					backgroundColor: "var(--vscode-badge-background)",
 					backgroundColor: "var(--vscode-badge-background)",
 					color: "var(--vscode-badge-foreground)",
 					color: "var(--vscode-badge-foreground)",
 					borderRadius: "3px",
 					borderRadius: "3px",
-					padding: "8px",
+					padding: "12px",
 					display: "flex",
 					display: "flex",
 					flexDirection: "column",
 					flexDirection: "column",
 					gap: "8px",
 					gap: "8px",
 				}}>
 				}}>
+				<div
+					style={{
+						display: "flex",
+						justifyContent: "space-between",
+						alignItems: "center",
+					}}>
+					<span style={{ fontWeight: "bold", fontSize: "16px" }}>Task</span>
+					<VSCodeButton
+						appearance="icon"
+						onClick={onClose}
+						style={{ marginTop: "-5px", marginRight: "-5px" }}>
+						<span className="codicon codicon-close"></span>
+					</VSCodeButton>
+				</div>
 				<div style={{ fontSize: "var(--vscode-font-size)", lineHeight: "1.5" }}>
 				<div style={{ fontSize: "var(--vscode-font-size)", lineHeight: "1.5" }}>
 					<TextTruncate
 					<TextTruncate
 						line={isExpanded ? 0 : 3}
 						line={isExpanded ? 0 : 3}
@@ -58,20 +71,24 @@ const TaskHeader: React.FC<TaskHeaderProps> = ({ taskText, tokensIn, tokensOut,
 					)}
 					)}
 				</div>
 				</div>
 				<div style={{ display: "flex", flexDirection: "column", gap: "4px" }}>
 				<div style={{ display: "flex", flexDirection: "column", gap: "4px" }}>
-					<div style={{ display: "flex", justifyContent: "space-between", alignItems: "center" }}>
+					<div style={{ display: "flex", alignItems: "center", gap: "4px" }}>
 						<span style={{ fontWeight: "bold" }}>Tokens:</span>
 						<span style={{ fontWeight: "bold" }}>Tokens:</span>
-						<div style={{ display: "flex", gap: "8px" }}>
-							<span style={{ display: "flex", alignItems: "center", gap: "4px" }}>
-								<i className="codicon codicon-arrow-down" style={{ fontSize: "12px", marginBottom: "-1px" }} />
-								{tokensIn.toLocaleString()}
-							</span>
-							<span style={{ display: "flex", alignItems: "center", gap: "4px" }}>
-								<i className="codicon codicon-arrow-up" style={{ fontSize: "12px", marginBottom: "-1px" }} />
-								{tokensOut.toLocaleString()}
-							</span>
-						</div>
+						<span style={{ display: "flex", alignItems: "center", gap: "4px" }}>
+							<i
+								className="codicon codicon-arrow-down"
+								style={{ fontSize: "12px", marginBottom: "-2px" }}
+							/>
+							{tokensIn.toLocaleString()}
+						</span>
+						<span style={{ display: "flex", alignItems: "center", gap: "4px" }}>
+							<i
+								className="codicon codicon-arrow-up"
+								style={{ fontSize: "12px", marginBottom: "-2px" }}
+							/>
+							{tokensOut.toLocaleString()}
+						</span>
 					</div>
 					</div>
-					<div style={{ display: "flex", justifyContent: "space-between", alignItems: "center" }}>
+					<div style={{ display: "flex", alignItems: "center", gap: "4px" }}>
 						<span style={{ fontWeight: "bold" }}>API Cost:</span>
 						<span style={{ fontWeight: "bold" }}>API Cost:</span>
 						<span>${totalCost.toFixed(4)}</span>
 						<span>${totalCost.toFixed(4)}</span>
 					</div>
 					</div>

+ 7 - 0
webview-ui/src/utilities/combineCommandSequences.ts

@@ -27,6 +27,7 @@ export function combineCommandSequences(messages: ClaudeMessage[]): ClaudeMessag
 	for (let i = 0; i < messages.length; i++) {
 	for (let i = 0; i < messages.length; i++) {
 		if (messages[i].type === "ask" && messages[i].ask === "command") {
 		if (messages[i].type === "ask" && messages[i].ask === "command") {
 			let combinedText = messages[i].text || ""
 			let combinedText = messages[i].text || ""
+			let didAddOutput = false
 			let j = i + 1
 			let j = i + 1
 
 
 			while (j < messages.length) {
 			while (j < messages.length) {
@@ -35,6 +36,11 @@ export function combineCommandSequences(messages: ClaudeMessage[]): ClaudeMessag
 					break
 					break
 				}
 				}
 				if (messages[j].type === "say" && messages[j].say === "command_output") {
 				if (messages[j].type === "say" && messages[j].say === "command_output") {
+					if (!didAddOutput) {
+						// Add a newline before the first output
+						combinedText += `\n${COMMAND_OUTPUT_STRING}`
+						didAddOutput = true
+					}
 					combinedText += "\n" + (messages[j].text || "")
 					combinedText += "\n" + (messages[j].text || "")
 				}
 				}
 				j++
 				j++
@@ -60,3 +66,4 @@ export function combineCommandSequences(messages: ClaudeMessage[]): ClaudeMessag
 			return msg
 			return msg
 		})
 		})
 }
 }
+export const COMMAND_OUTPUT_STRING = "Output:"