Browse Source

Split api and chat message persistence into a separate module (#2866)

Chris Estreich 10 months ago
parent
commit
72962b3c76

+ 22 - 0
package-lock.json

@@ -40,6 +40,7 @@
 				"js-tiktoken": "^1.0.19",
 				"js-tiktoken": "^1.0.19",
 				"mammoth": "^1.8.0",
 				"mammoth": "^1.8.0",
 				"monaco-vscode-textmate-theme-converter": "^0.1.7",
 				"monaco-vscode-textmate-theme-converter": "^0.1.7",
+				"node-cache": "^5.1.2",
 				"node-ipc": "^12.0.0",
 				"node-ipc": "^12.0.0",
 				"openai": "^4.78.1",
 				"openai": "^4.78.1",
 				"os-name": "^6.0.0",
 				"os-name": "^6.0.0",
@@ -16784,6 +16785,27 @@
 			"license": "MIT",
 			"license": "MIT",
 			"optional": true
 			"optional": true
 		},
 		},
+		"node_modules/node-cache": {
+			"version": "5.1.2",
+			"resolved": "https://registry.npmjs.org/node-cache/-/node-cache-5.1.2.tgz",
+			"integrity": "sha512-t1QzWwnk4sjLWaQAS8CHgOJ+RAfmHpxFWmc36IWTiWHQfs0w5JDMBS1b1ZxQteo0vVVuWJvIUKHDkkeK7vIGCg==",
+			"license": "MIT",
+			"dependencies": {
+				"clone": "2.x"
+			},
+			"engines": {
+				"node": ">= 8.0.0"
+			}
+		},
+		"node_modules/node-cache/node_modules/clone": {
+			"version": "2.1.2",
+			"resolved": "https://registry.npmjs.org/clone/-/clone-2.1.2.tgz",
+			"integrity": "sha512-3Pe/CF1Nn94hyhIYpjtiLhdCoEoz0DqQ+988E9gmeEdQZlojxnOb74wctFyuwWQHzqyf9X7C7MG8juUpqBJT8w==",
+			"license": "MIT",
+			"engines": {
+				"node": ">=0.8"
+			}
+		},
 		"node_modules/node-domexception": {
 		"node_modules/node-domexception": {
 			"version": "1.0.0",
 			"version": "1.0.0",
 			"resolved": "https://registry.npmjs.org/node-domexception/-/node-domexception-1.0.0.tgz",
 			"resolved": "https://registry.npmjs.org/node-domexception/-/node-domexception-1.0.0.tgz",

+ 1 - 0
package.json

@@ -432,6 +432,7 @@
 		"js-tiktoken": "^1.0.19",
 		"js-tiktoken": "^1.0.19",
 		"mammoth": "^1.8.0",
 		"mammoth": "^1.8.0",
 		"monaco-vscode-textmate-theme-converter": "^0.1.7",
 		"monaco-vscode-textmate-theme-converter": "^0.1.7",
+		"node-cache": "^5.1.2",
 		"node-ipc": "^12.0.0",
 		"node-ipc": "^12.0.0",
 		"openai": "^4.78.1",
 		"openai": "^4.78.1",
 		"os-name": "^6.0.0",
 		"os-name": "^6.0.0",

+ 24 - 82
src/core/Cline.ts

@@ -1,4 +1,3 @@
-import fs from "fs/promises"
 import * as path from "path"
 import * as path from "path"
 import os from "os"
 import os from "os"
 import crypto from "crypto"
 import crypto from "crypto"
@@ -8,7 +7,6 @@ import { Anthropic } from "@anthropic-ai/sdk"
 import cloneDeep from "clone-deep"
 import cloneDeep from "clone-deep"
 import delay from "delay"
 import delay from "delay"
 import pWaitFor from "p-wait-for"
 import pWaitFor from "p-wait-for"
-import getFolderSize from "get-folder-size"
 import { serializeError } from "serialize-error"
 import { serializeError } from "serialize-error"
 import * as vscode from "vscode"
 import * as vscode from "vscode"
 
 
@@ -58,7 +56,6 @@ import { TerminalRegistry } from "../integrations/terminal/TerminalRegistry"
 
 
 // utils
 // utils
 import { calculateApiCostAnthropic } from "../utils/cost"
 import { calculateApiCostAnthropic } from "../utils/cost"
-import { fileExistsAtPath } from "../utils/fs"
 import { arePathsEqual, getWorkspacePath } from "../utils/path"
 import { arePathsEqual, getWorkspacePath } from "../utils/path"
 
 
 // tools
 // tools
@@ -93,6 +90,7 @@ import { truncateConversationIfNeeded } from "./sliding-window"
 import { ClineProvider } from "./webview/ClineProvider"
 import { ClineProvider } from "./webview/ClineProvider"
 import { validateToolUse } from "./mode-validator"
 import { validateToolUse } from "./mode-validator"
 import { MultiSearchReplaceDiffStrategy } from "./diff/strategies/multi-search-replace"
 import { MultiSearchReplaceDiffStrategy } from "./diff/strategies/multi-search-replace"
+import { readApiMessages, saveApiMessages, readTaskMessages, saveTaskMessages, taskMetadata } from "./task-persistence"
 
 
 type UserContent = Array<Anthropic.Messages.ContentBlockParam>
 type UserContent = Array<Anthropic.Messages.ContentBlockParam>
 
 
@@ -171,6 +169,7 @@ export class Cline extends EventEmitter<ClineEvents> {
 
 
 	// Not private since it needs to be accessible by tools.
 	// Not private since it needs to be accessible by tools.
 	providerRef: WeakRef<ClineProvider>
 	providerRef: WeakRef<ClineProvider>
+	private readonly globalStoragePath: string
 	private abort: boolean = false
 	private abort: boolean = false
 	didFinishAbortingStream = false
 	didFinishAbortingStream = false
 	abandoned = false
 	abandoned = false
@@ -244,6 +243,7 @@ export class Cline extends EventEmitter<ClineEvents> {
 		this.fuzzyMatchThreshold = fuzzyMatchThreshold
 		this.fuzzyMatchThreshold = fuzzyMatchThreshold
 		this.consecutiveMistakeLimit = consecutiveMistakeLimit
 		this.consecutiveMistakeLimit = consecutiveMistakeLimit
 		this.providerRef = new WeakRef(provider)
 		this.providerRef = new WeakRef(provider)
+		this.globalStoragePath = provider.context.globalStorageUri.fsPath
 		this.diffViewProvider = new DiffViewProvider(this.cwd)
 		this.diffViewProvider = new DiffViewProvider(this.cwd)
 		this.enableCheckpoints = enableCheckpoints
 		this.enableCheckpoints = enableCheckpoints
 
 
@@ -294,26 +294,8 @@ export class Cline extends EventEmitter<ClineEvents> {
 
 
 	// Storing task to disk for history
 	// Storing task to disk for history
 
 
-	private async ensureTaskDirectoryExists(): Promise<string> {
-		const globalStoragePath = this.providerRef.deref()?.context.globalStorageUri.fsPath
-		if (!globalStoragePath) {
-			throw new Error("Global storage uri is invalid")
-		}
-
-		// Use storagePathManager to retrieve the task storage directory
-		const { getTaskDirectoryPath } = await import("../shared/storagePathManager")
-		return getTaskDirectoryPath(globalStoragePath, this.taskId)
-	}
-
 	private async getSavedApiConversationHistory(): Promise<Anthropic.MessageParam[]> {
 	private async getSavedApiConversationHistory(): Promise<Anthropic.MessageParam[]> {
-		const filePath = path.join(await this.ensureTaskDirectoryExists(), GlobalFileNames.apiConversationHistory)
-		const fileExists = await fileExistsAtPath(filePath)
-
-		if (fileExists) {
-			return JSON.parse(await fs.readFile(filePath, "utf8"))
-		}
-
-		return []
+		return readApiMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath })
 	}
 	}
 
 
 	private async addToApiConversationHistory(message: Anthropic.MessageParam) {
 	private async addToApiConversationHistory(message: Anthropic.MessageParam) {
@@ -329,8 +311,11 @@ export class Cline extends EventEmitter<ClineEvents> {
 
 
 	private async saveApiConversationHistory() {
 	private async saveApiConversationHistory() {
 		try {
 		try {
-			const filePath = path.join(await this.ensureTaskDirectoryExists(), GlobalFileNames.apiConversationHistory)
-			await fs.writeFile(filePath, JSON.stringify(this.apiConversationHistory))
+			await saveApiMessages({
+				messages: this.apiConversationHistory,
+				taskId: this.taskId,
+				globalStoragePath: this.globalStoragePath,
+			})
 		} catch (error) {
 		} catch (error) {
 			// in the off chance this fails, we don't want to stop the task
 			// in the off chance this fails, we don't want to stop the task
 			console.error("Failed to save API conversation history:", error)
 			console.error("Failed to save API conversation history:", error)
@@ -338,21 +323,7 @@ export class Cline extends EventEmitter<ClineEvents> {
 	}
 	}
 
 
 	private async getSavedClineMessages(): Promise<ClineMessage[]> {
 	private async getSavedClineMessages(): Promise<ClineMessage[]> {
-		const filePath = path.join(await this.ensureTaskDirectoryExists(), GlobalFileNames.uiMessages)
-
-		if (await fileExistsAtPath(filePath)) {
-			return JSON.parse(await fs.readFile(filePath, "utf8"))
-		} else {
-			// check old location
-			const oldPath = path.join(await this.ensureTaskDirectoryExists(), "claude_messages.json")
-			if (await fileExistsAtPath(oldPath)) {
-				const data = JSON.parse(await fs.readFile(oldPath, "utf8"))
-				await fs.unlink(oldPath) // remove old file
-				return data
-			}
-		}
-
-		return []
+		return readTaskMessages({ taskId: this.taskId, globalStoragePath: this.globalStoragePath })
 	}
 	}
 
 
 	private async addToClineMessages(message: ClineMessage) {
 	private async addToClineMessages(message: ClineMessage) {
@@ -374,54 +345,25 @@ export class Cline extends EventEmitter<ClineEvents> {
 		this.emit("message", { action: "updated", message: partialMessage })
 		this.emit("message", { action: "updated", message: partialMessage })
 	}
 	}
 
 
-	private taskDirSize = 0
-	private taskDirSizeCheckedAt = 0
-	private readonly taskDirSizeCheckInterval = 30_000
-
 	private async saveClineMessages() {
 	private async saveClineMessages() {
 		try {
 		try {
-			const taskDir = await this.ensureTaskDirectoryExists()
-			const filePath = path.join(taskDir, GlobalFileNames.uiMessages)
-			await fs.writeFile(filePath, JSON.stringify(this.clineMessages))
-
-			const tokenUsage = this.getTokenUsage()
-			this.emit("taskTokenUsageUpdated", this.taskId, tokenUsage)
-
-			const taskMessage = this.clineMessages[0] // First message is always the task say
-
-			const lastRelevantMessage =
-				this.clineMessages[
-					findLastIndex(
-						this.clineMessages,
-						(m) => !(m.ask === "resume_task" || m.ask === "resume_completed_task"),
-					)
-				]
-
-			if (Date.now() - this.taskDirSizeCheckedAt > this.taskDirSizeCheckInterval) {
-				this.taskDirSizeCheckedAt = Date.now()
-
-				try {
-					this.taskDirSize = await getFolderSize.loose(taskDir)
-				} catch (err) {
-					console.error(
-						`[saveClineMessages] failed to get task directory size (${taskDir}): ${err instanceof Error ? err.message : String(err)}`,
-					)
-				}
-			}
+			await saveTaskMessages({
+				messages: this.clineMessages,
+				taskId: this.taskId,
+				globalStoragePath: this.globalStoragePath,
+			})
 
 
-			await this.providerRef.deref()?.updateTaskHistory({
-				id: this.taskId,
-				number: this.taskNumber,
-				ts: lastRelevantMessage.ts,
-				task: taskMessage.text ?? "",
-				tokensIn: tokenUsage.totalTokensIn,
-				tokensOut: tokenUsage.totalTokensOut,
-				cacheWrites: tokenUsage.totalCacheWrites,
-				cacheReads: tokenUsage.totalCacheReads,
-				totalCost: tokenUsage.totalCost,
-				size: this.taskDirSize,
+			const { historyItem, tokenUsage } = await taskMetadata({
+				messages: this.clineMessages,
+				taskId: this.taskId,
+				taskNumber: this.taskNumber,
+				globalStoragePath: this.globalStoragePath,
 				workspace: this.cwd,
 				workspace: this.cwd,
 			})
 			})
+
+			this.emit("taskTokenUsageUpdated", this.taskId, tokenUsage)
+
+			await this.providerRef.deref()?.updateTaskHistory(historyItem)
 		} catch (error) {
 		} catch (error) {
 			console.error("Failed to save cline messages:", error)
 			console.error("Failed to save cline messages:", error)
 		}
 		}

+ 50 - 0
src/core/task-persistence/apiMessages.ts

@@ -0,0 +1,50 @@
+import * as path from "path"
+import * as fs from "fs/promises"
+
+import { Anthropic } from "@anthropic-ai/sdk"
+
+import { fileExistsAtPath } from "../../utils/fs"
+
+import { GlobalFileNames } from "../../shared/globalFileNames"
+import { getTaskDirectoryPath } from "../../shared/storagePathManager"
+
+export type ApiMessage = Anthropic.MessageParam & { ts?: number }
+
+export async function readApiMessages({
+	taskId,
+	globalStoragePath,
+}: {
+	taskId: string
+	globalStoragePath: string
+}): Promise<ApiMessage[]> {
+	const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId)
+	const filePath = path.join(taskDir, GlobalFileNames.apiConversationHistory)
+
+	if (await fileExistsAtPath(filePath)) {
+		return JSON.parse(await fs.readFile(filePath, "utf8"))
+	} else {
+		const oldPath = path.join(taskDir, "claude_messages.json")
+
+		if (await fileExistsAtPath(oldPath)) {
+			const data = JSON.parse(await fs.readFile(oldPath, "utf8"))
+			await fs.unlink(oldPath)
+			return data
+		}
+	}
+
+	return []
+}
+
+export async function saveApiMessages({
+	messages,
+	taskId,
+	globalStoragePath,
+}: {
+	messages: ApiMessage[]
+	taskId: string
+	globalStoragePath: string
+}) {
+	const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId)
+	const filePath = path.join(taskDir, GlobalFileNames.apiConversationHistory)
+	await fs.writeFile(filePath, JSON.stringify(messages))
+}

+ 3 - 0
src/core/task-persistence/index.ts

@@ -0,0 +1,3 @@
+export { readApiMessages, saveApiMessages } from "./apiMessages"
+export { readTaskMessages, saveTaskMessages } from "./taskMessages"
+export { taskMetadata } from "./taskMetadata"

+ 40 - 0
src/core/task-persistence/taskMessages.ts

@@ -0,0 +1,40 @@
+import * as path from "path"
+import * as fs from "fs/promises"
+
+import { fileExistsAtPath } from "../../utils/fs"
+
+import { GlobalFileNames } from "../../shared/globalFileNames"
+import { ClineMessage } from "../../shared/ExtensionMessage"
+import { getTaskDirectoryPath } from "../../shared/storagePathManager"
+
+export type ReadTaskMessagesOptions = {
+	taskId: string
+	globalStoragePath: string
+}
+
+export async function readTaskMessages({
+	taskId,
+	globalStoragePath,
+}: ReadTaskMessagesOptions): Promise<ClineMessage[]> {
+	const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId)
+	const filePath = path.join(taskDir, GlobalFileNames.uiMessages)
+	const fileExists = await fileExistsAtPath(filePath)
+
+	if (fileExists) {
+		return JSON.parse(await fs.readFile(filePath, "utf8"))
+	}
+
+	return []
+}
+
+export type SaveTaskMessagesOptions = {
+	messages: ClineMessage[]
+	taskId: string
+	globalStoragePath: string
+}
+
+export async function saveTaskMessages({ messages, taskId, globalStoragePath }: SaveTaskMessagesOptions) {
+	const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId)
+	const filePath = path.join(taskDir, GlobalFileNames.uiMessages)
+	await fs.writeFile(filePath, JSON.stringify(messages))
+}

+ 63 - 0
src/core/task-persistence/taskMetadata.ts

@@ -0,0 +1,63 @@
+import NodeCache from "node-cache"
+import getFolderSize from "get-folder-size"
+
+import { ClineMessage } from "../../shared/ExtensionMessage"
+import { combineApiRequests } from "../../shared/combineApiRequests"
+import { combineCommandSequences } from "../../shared/combineCommandSequences"
+import { getApiMetrics } from "../../shared/getApiMetrics"
+import { findLastIndex } from "../../shared/array"
+import { HistoryItem } from "../../shared/HistoryItem"
+import { getTaskDirectoryPath } from "../../shared/storagePathManager"
+
+const taskSizeCache = new NodeCache({ stdTTL: 5 * 60, checkperiod: 30 })
+
+export type TaskMetadataOptions = {
+	messages: ClineMessage[]
+	taskId: string
+	taskNumber: number
+	globalStoragePath: string
+	workspace: string
+}
+
+export async function taskMetadata({
+	messages,
+	taskId,
+	taskNumber,
+	globalStoragePath,
+	workspace,
+}: TaskMetadataOptions) {
+	const taskDir = await getTaskDirectoryPath(globalStoragePath, taskId)
+	const taskMessage = messages[0] // First message is always the task say.
+
+	const lastRelevantMessage =
+		messages[findLastIndex(messages, (m) => !(m.ask === "resume_task" || m.ask === "resume_completed_task"))]
+
+	let taskDirSize = taskSizeCache.get<number>(taskDir)
+
+	if (taskDirSize === undefined) {
+		try {
+			taskDirSize = await getFolderSize.loose(taskDir)
+			taskSizeCache.set<number>(taskDir, taskDirSize)
+		} catch (error) {
+			taskDirSize = 0
+		}
+	}
+
+	const tokenUsage = getApiMetrics(combineApiRequests(combineCommandSequences(messages.slice(1))))
+
+	const historyItem: HistoryItem = {
+		id: taskId,
+		number: taskNumber,
+		ts: lastRelevantMessage.ts,
+		task: taskMessage.text ?? "",
+		tokensIn: tokenUsage.totalTokensIn,
+		tokensOut: tokenUsage.totalTokensOut,
+		cacheWrites: tokenUsage.totalCacheWrites,
+		cacheReads: tokenUsage.totalCacheReads,
+		totalCost: tokenUsage.totalCost,
+		size: taskDirSize,
+		workspace,
+	}
+
+	return { historyItem, tokenUsage }
+}