move task ephemeral state to state class

src/api/retry.ts

@@ -57,7 +57,7 @@ export function withRetry(options: RetryOptions = {}) {
 					const handlerInstance = this as any
 					if (handlerInstance.options?.onRetryAttempt) {
 						try {
-							handlerInstance.options.onRetryAttempt(attempt + 1, maxRetries, delay, error)
+							await handlerInstance.options.onRetryAttempt(attempt + 1, maxRetries, delay, error)
 						} catch (e) {
 							console.error("Error in onRetryAttempt callback:", e)
 						}

src/core/controller/checkpoints/checkpointRestore.ts

@@ -9,7 +9,7 @@ export async function checkpointRestore(controller: Controller, request: Checkpo
 
 	if (request.number) {
 		// wait for messages to be loaded
-		await pWaitFor(() => controller.task?.isInitialized === true, {
+		await pWaitFor(() => controller.task?.taskState.isInitialized === true, {
 			timeout: 3_000,
 		}).catch(() => {
 			console.error("Failed to init new cline instance")

src/core/controller/index.ts

@@ -445,8 +445,8 @@ export class Controller {
 
 		if (this.task) {
 			this.task.chatSettings = chatSettings
-			if (this.task.isAwaitingPlanResponse && didSwitchToActMode) {
-				this.task.didRespondToPlanAskBySwitchingMode = true
+			if (this.task.taskState.isAwaitingPlanResponse && didSwitchToActMode) {
+				this.task.taskState.didRespondToPlanAskBySwitchingMode = true
 				// Use chatContent if provided, otherwise use default message
 				await this.task.handleWebviewAskResponse(
 					"messageResponse",
@@ -471,9 +471,9 @@ export class Controller {
 			await pWaitFor(
 				() =>
 					this.task === undefined ||
-					this.task.isStreaming === false ||
-					this.task.didFinishAbortingStream ||
-					this.task.isWaitingForFirstChunk, // if only first chunk is processed, then there's no need to wait for graceful abort (closes edits, browser, etc)
+					this.task.taskState.isStreaming === false ||
+					this.task.taskState.didFinishAbortingStream ||
+					this.task.taskState.isWaitingForFirstChunk, // if only first chunk is processed, then there's no need to wait for graceful abort (closes edits, browser, etc)
 				{
 					timeout: 3_000,
 				},
@@ -482,7 +482,7 @@ export class Controller {
 			})
 			if (this.task) {
 				// 'abandoned' will prevent this cline instance from affecting future cline instance gui. this may happen if its hanging on a streaming request
-				this.task.abandoned = true
+				this.task.taskState.abandoned = true
 			}
 			await this.initTask(undefined, undefined, undefined, historyItem) // clears task again, so we need to abortTask manually above
 			// await this.postStateToWebview() // new Cline instance will post state when it's ready. having this here sent an empty messages array to webview leading to virtuoso having to reload the entire list
@@ -1006,8 +1006,8 @@ export class Controller {
 			apiConfiguration,
 			uriScheme: vscode.env.uriScheme,
 			currentTaskItem: this.task?.taskId ? (taskHistory || []).find((item) => item.id === this.task?.taskId) : undefined,
-			checkpointTrackerErrorMessage: this.task?.checkpointTrackerErrorMessage,
-			clineMessages: this.task?.clineMessages || [],
+			checkpointTrackerErrorMessage: this.task?.taskState.checkpointTrackerErrorMessage,
+			clineMessages: this.task?.messageStateHandler.getClineMessages() || [],
 			taskHistory: (taskHistory || [])
 				.filter((item) => item.ts && item.task)
 				.sort((a, b) => b.ts - a.ts)

src/core/task/TaskState.ts

@@ -0,0 +1,57 @@
+import { Anthropic } from "@anthropic-ai/sdk"
+import { AssistantMessageContent } from "@core/assistant-message"
+import { StreamingJsonReplacer } from "@core/assistant-message/diff-json"
+import { ClineAskResponse } from "@shared/WebviewMessage"
+
+export class TaskState {
+	// Streaming flags
+	isStreaming = false
+	isWaitingForFirstChunk = false
+	didCompleteReadingStream = false
+
+	// Content processing
+	currentStreamingContentIndex = 0
+	assistantMessageContent: AssistantMessageContent[] = []
+	userMessageContent: (Anthropic.TextBlockParam | Anthropic.ImageBlockParam)[] = []
+	userMessageContentReady = false
+
+	// Presentation locks
+	presentAssistantMessageLocked = false
+	presentAssistantMessageHasPendingUpdates = false
+
+	// Claude 4 experimental JSON streaming
+	streamingJsonReplacer?: StreamingJsonReplacer
+	lastProcessedJsonLength: number = 0
+
+	// Ask/Response handling
+	askResponse?: ClineAskResponse
+	askResponseText?: string
+	askResponseImages?: string[]
+	askResponseFiles?: string[]
+	lastMessageTs?: number
+
+	// Plan mode specific state
+	isAwaitingPlanResponse = false
+	didRespondToPlanAskBySwitchingMode = false
+
+	// Tool execution flags
+	didRejectTool = false
+	didAlreadyUseTool = false
+	didEditFile: boolean = false
+
+	// Consecutive request tracking
+	consecutiveAutoApprovedRequestsCount: number = 0
+
+	// Error tracking
+	consecutiveMistakeCount: number = 0
+	didAutomaticallyRetryFailedApiRequest = false
+	checkpointTrackerErrorMessage?: string
+
+	// Task Initialization
+	isInitialized = false
+
+	// Task Abort / Cancellation
+	abort: boolean = false
+	didFinishAbortingStream = false
+	abandoned = false
+}

src/core/task/message-state.ts

@@ -1,5 +1,5 @@
 import { combineApiRequests } from "@/shared/combineApiRequests"
-import { ensureTaskDirectoryExists, saveClineMessages } from "../storage/disk"
+import { ensureTaskDirectoryExists, saveApiConversationHistory, saveClineMessages } from "../storage/disk"
 import * as vscode from "vscode"
 import { ClineMessage } from "@/shared/ExtensionMessage"
 import { getApiMetrics } from "@/shared/getApiMetrics"
@@ -10,57 +10,132 @@ import os from "os"
 import * as path from "path"
 import CheckpointTracker from "@integrations/checkpoints/CheckpointTracker"
 import { HistoryItem } from "@/shared/HistoryItem"
+import Anthropic from "@anthropic-ai/sdk"
+
+interface MessageStateHandlerParams {
+	context: vscode.ExtensionContext
+	taskId: string
+	conversationHistoryDeletedRange?: [number, number]
+	taskIsFavorited?: boolean
+	updateTaskHistory: (historyItem: HistoryItem) => Promise<HistoryItem[]>
+}
 
 const cwd = vscode.workspace.workspaceFolders?.map((folder) => folder.uri.fsPath).at(0) ?? path.join(os.homedir(), "Desktop") // may or may not exist but fs checking existence would immediately ask for permission which would be bad UX, need to come up with a better solution
 
-// need to call getContext() from the task object when passing in context
-export async function saveClineMessagesAndUpdateHistory(
-	context: vscode.ExtensionContext,
-	taskId: string,
-	clineMessages: ClineMessage[],
-	taskIsFavorited: boolean,
-	conversationHistoryDeletedRange: [number, number] | undefined,
-	checkpointTracker: CheckpointTracker | undefined,
-	updateTaskHistory: (historyItem: HistoryItem) => Promise<HistoryItem[]>,
-) {
-	try {
-		await saveClineMessages(context, taskId, clineMessages)
+export class MessageStateHandler {
+	private apiConversationHistory: Anthropic.MessageParam[] = []
+	private clineMessages: ClineMessage[] = []
+	private conversationHistoryDeletedRange: [number, number] | undefined
+	private taskIsFavorited: boolean
+	private checkpointTracker: CheckpointTracker | undefined
+	private updateTaskHistory: (historyItem: HistoryItem) => Promise<HistoryItem[]>
+	private context: vscode.ExtensionContext
+	private taskId: string
 
-		// combined as they are in ChatView
-		const apiMetrics = getApiMetrics(combineApiRequests(combineCommandSequences(clineMessages.slice(1))))
-		const taskMessage = clineMessages[0] // first message is always the task say
-		const lastRelevantMessage =
-			clineMessages[
-				findLastIndex(
-					clineMessages,
-					(message) => !(message.ask === "resume_task" || message.ask === "resume_completed_task"),
-				)
-			]
-		const taskDir = await ensureTaskDirectoryExists(context, taskId)
-		let taskDirSize = 0
+	constructor(params: MessageStateHandlerParams) {
+		this.context = params.context
+		this.taskId = params.taskId
+		this.conversationHistoryDeletedRange = params.conversationHistoryDeletedRange
+		this.taskIsFavorited = params.taskIsFavorited ?? false
+		this.updateTaskHistory = params.updateTaskHistory
+	}
+
+	setCheckpointTracker(tracker: CheckpointTracker | undefined) {
+		this.checkpointTracker = tracker
+	}
+
+	getApiConversationHistory(): Anthropic.MessageParam[] {
+		return this.apiConversationHistory
+	}
+
+	setApiConversationHistory(newHistory: Anthropic.MessageParam[]): void {
+		this.apiConversationHistory = newHistory
+	}
+
+	getClineMessages(): ClineMessage[] {
+		return this.clineMessages
+	}
+
+	setClineMessages(newMessages: ClineMessage[]) {
+		this.clineMessages = newMessages
+	}
+
+	async saveClineMessagesAndUpdateHistory(): Promise<void> {
 		try {
-			// getFolderSize.loose silently ignores errors
-			// returns # of bytes, size/1000/1000 = MB
-			taskDirSize = await getFolderSize.loose(taskDir)
+			await saveClineMessages(this.context, this.taskId, this.clineMessages)
+
+			// combined as they are in ChatView
+			const apiMetrics = getApiMetrics(combineApiRequests(combineCommandSequences(this.clineMessages.slice(1))))
+			const taskMessage = this.clineMessages[0] // first message is always the task say
+			const lastRelevantMessage =
+				this.clineMessages[
+					findLastIndex(
+						this.clineMessages,
+						(message) => !(message.ask === "resume_task" || message.ask === "resume_completed_task"),
+					)
+				]
+			const taskDir = await ensureTaskDirectoryExists(this.context, this.taskId)
+			let taskDirSize = 0
+			try {
+				// getFolderSize.loose silently ignores errors
+				// returns # of bytes, size/1000/1000 = MB
+				taskDirSize = await getFolderSize.loose(taskDir)
+			} catch (error) {
+				console.error("Failed to get task directory size:", taskDir, error)
+			}
+			await this.updateTaskHistory({
+				id: this.taskId,
+				ts: lastRelevantMessage.ts,
+				task: taskMessage.text ?? "",
+				tokensIn: apiMetrics.totalTokensIn,
+				tokensOut: apiMetrics.totalTokensOut,
+				cacheWrites: apiMetrics.totalCacheWrites,
+				cacheReads: apiMetrics.totalCacheReads,
+				totalCost: apiMetrics.totalCost,
+				size: taskDirSize,
+				shadowGitConfigWorkTree: await this.checkpointTracker?.getShadowGitConfigWorkTree(),
+				cwdOnTaskInitialization: cwd,
+				conversationHistoryDeletedRange: this.conversationHistoryDeletedRange,
+				isFavorited: this.taskIsFavorited,
+			})
 		} catch (error) {
-			console.error("Failed to get task directory size:", taskDir, error)
+			console.error("Failed to save cline messages:", error)
 		}
-		await updateTaskHistory({
-			id: taskId,
-			ts: lastRelevantMessage.ts,
-			task: taskMessage.text ?? "",
-			tokensIn: apiMetrics.totalTokensIn,
-			tokensOut: apiMetrics.totalTokensOut,
-			cacheWrites: apiMetrics.totalCacheWrites,
-			cacheReads: apiMetrics.totalCacheReads,
-			totalCost: apiMetrics.totalCost,
-			size: taskDirSize,
-			shadowGitConfigWorkTree: await checkpointTracker?.getShadowGitConfigWorkTree(),
-			cwdOnTaskInitialization: cwd,
-			conversationHistoryDeletedRange: conversationHistoryDeletedRange,
-			isFavorited: taskIsFavorited,
-		})
-	} catch (error) {
-		console.error("Failed to save cline messages:", error)
+	}
+
+	async addToApiConversationHistory(message: Anthropic.MessageParam) {
+		this.apiConversationHistory.push(message)
+		await saveApiConversationHistory(this.context, this.taskId, this.apiConversationHistory)
+	}
+
+	async overwriteApiConversationHistory(newHistory: Anthropic.MessageParam[]): Promise<void> {
+		this.apiConversationHistory = newHistory
+		await saveApiConversationHistory(this.context, this.taskId, this.apiConversationHistory)
+	}
+
+	async addToClineMessages(message: ClineMessage) {
+		// these values allow us to reconstruct the conversation history at the time this cline message was created
+		// it's important that apiConversationHistory is initialized before we add cline messages
+		message.conversationHistoryIndex = this.apiConversationHistory.length - 1 // NOTE: this is the index of the last added message which is the user message, and once the clinemessages have been presented we update the apiconversationhistory with the completed assistant message. This means when resetting to a message, we need to +1 this index to get the correct assistant message that this tool use corresponds to
+		message.conversationHistoryDeletedRange = this.conversationHistoryDeletedRange
+		this.clineMessages.push(message)
+		await this.saveClineMessagesAndUpdateHistory()
+	}
+
+	async overwriteClineMessages(newMessages: ClineMessage[]) {
+		this.clineMessages = newMessages
+		await this.saveClineMessagesAndUpdateHistory()
+	}
+
+	async updateClineMessage(index: number, updates: Partial<ClineMessage>): Promise<void> {
+		if (index < 0 || index >= this.clineMessages.length) {
+			throw new Error(`Invalid message index: ${index}`)
+		}
+
+		// Apply updates to the message
+		Object.assign(this.clineMessages[index], updates)
+
+		// Save changes and update history
+		await this.saveClineMessagesAndUpdateHistory()
 	}
 }

src/core/task/utils.ts

@@ -0,0 +1,70 @@
+import { showSystemNotification } from "@/integrations/notifications"
+import { ClineApiReqCancelReason, ClineApiReqInfo } from "@/shared/ExtensionMessage"
+import { serializeError } from "serialize-error"
+import { MessageStateHandler } from "./message-state"
+import { calculateApiCostAnthropic } from "@/utils/cost"
+import { ApiHandler } from "@/api"
+
+export function formatErrorWithStatusCode(error: any): string {
+	const statusCode = error.status || error.statusCode || (error.response && error.response.status)
+	const message = error.message ?? JSON.stringify(serializeError(error), null, 2)
+
+	// Only prepend the statusCode if it's not already part of the message
+	return statusCode && !message.includes(statusCode.toString()) ? `${statusCode} - ${message}` : message
+}
+
+export const showNotificationForApprovalIfAutoApprovalEnabled = (
+	message: string,
+	autoApprovalSettingsEnabled: boolean,
+	notificationsEnabled: boolean,
+) => {
+	if (autoApprovalSettingsEnabled && notificationsEnabled) {
+		showSystemNotification({
+			subtitle: "Approval Required",
+			message,
+		})
+	}
+}
+
+type UpdateApiReqMsgParams = {
+	messageStateHandler: MessageStateHandler
+	lastApiReqIndex: number
+	inputTokens: number
+	outputTokens: number
+	cacheWriteTokens: number
+	cacheReadTokens: number
+	totalCost?: number
+	api: ApiHandler
+	cancelReason?: ClineApiReqCancelReason
+	streamingFailedMessage?: string
+}
+
+// update api_req_started. we can't use api_req_finished anymore since it's a unique case where it could come after a streaming message (ie in the middle of being updated or executed)
+// fortunately api_req_finished was always parsed out for the gui anyways, so it remains solely for legacy purposes to keep track of prices in tasks from history
+// (it's worth removing a few months from now)
+export const updateApiReqMsg = async (params: UpdateApiReqMsgParams) => {
+	const clineMessages = params.messageStateHandler.getClineMessages()
+	const currentApiReqInfo: ClineApiReqInfo = JSON.parse(clineMessages[params.lastApiReqIndex].text || "{}")
+	delete currentApiReqInfo.retryStatus // Clear retry status when request is finalized
+
+	await params.messageStateHandler.updateClineMessage(params.lastApiReqIndex, {
+		text: JSON.stringify({
+			...currentApiReqInfo, // Spread the modified info (with retryStatus removed)
+			tokensIn: params.inputTokens,
+			tokensOut: params.outputTokens,
+			cacheWrites: params.cacheWriteTokens,
+			cacheReads: params.cacheReadTokens,
+			cost:
+				params.totalCost ??
+				calculateApiCostAnthropic(
+					params.api.getModel().info,
+					params.inputTokens,
+					params.outputTokens,
+					params.cacheWriteTokens,
+					params.cacheReadTokens,
+				),
+			cancelReason: params.cancelReason,
+			streamingFailedMessage: params.streamingFailedMessage,
+		} satisfies ClineApiReqInfo),
+	})
+}