diff --git a/apps/cli/package.json b/apps/cli/package.json index 73420e8a2d..f580abfca6 100644 --- a/apps/cli/package.json +++ b/apps/cli/package.json @@ -20,6 +20,7 @@ }, "dependencies": { "@inkjs/ui": "^2.0.0", + "@roo-code/core": "workspace:^", "@roo-code/types": "workspace:^", "@roo-code/vscode-shim": "workspace:^", "@vscode/ripgrep": "^1.15.9", diff --git a/apps/cli/src/ui/App.tsx b/apps/cli/src/ui/App.tsx index fe0fc38633..1bec82ddc0 100644 --- a/apps/cli/src/ui/App.tsx +++ b/apps/cli/src/ui/App.tsx @@ -4,7 +4,12 @@ import { useState, useEffect, useCallback, useRef, useMemo } from "react" import { EventEmitter } from "events" import { randomUUID } from "crypto" +// Import only message-utils to avoid custom-tools dependencies (execa/child_process) +import { consolidateTokenUsage, consolidateApiRequests, consolidateCommands } from "@roo-code/core/message-utils" +import type { ClineMessage } from "@roo-code/types" + import { useCLIStore } from "./store.js" +import { getContextWindow } from "../utils/getContextWindow.js" import Header from "./components/Header.js" import ChatHistoryItem from "./components/ChatHistoryItem.js" import LoadingText from "./components/LoadingText.js" @@ -166,8 +171,20 @@ function AppInner({ allSlashCommands, setFileSearchResults, setAllSlashCommands, + tokenUsage, + routerModels, + apiConfiguration, + setTokenUsage, + setRouterModels, + setApiConfiguration, } = useCLIStore() + // Compute context window from router models and API configuration + const contextWindow = useMemo( + () => getContextWindow(routerModels, apiConfiguration), + [routerModels, apiConfiguration], + ) + const hostRef = useRef(null) // eslint-disable-next-line @typescript-eslint/no-explicit-any const autocompleteRef = useRef>(null) @@ -527,6 +544,16 @@ function AppInner({ handleAskMessage(ts, ask, text, partial) } } + + // Compute token usage metrics from clineMessages + // Skip first message (task prompt) as per webview UI pattern + if (clineMessages.length > 1) { + const processed = consolidateApiRequests( + consolidateCommands(clineMessages.slice(1) as ClineMessage[]), + ) + const metrics = consolidateTokenUsage(processed) + setTokenUsage(metrics) + } } } else if (msg.type === "messageUpdated") { const clineMessage = msg.clineMessage as Record @@ -562,9 +589,29 @@ function AppInner({ source: cmd.source, })) setAllSlashCommands(slashCommands) + } else if (msg.type === "routerModels") { + // Handle router models for context window lookup + const models = msg.models as Record> | undefined + if (models) { + setRouterModels(models) + } + } else if (msg.type === "apiConfiguration") { + // Handle API configuration for model identification + const config = msg.configuration as unknown + if (config) { + setApiConfiguration(config as import("@roo-code/types").ProviderSettings) + } } }, - [handleSayMessage, handleAskMessage, setFileSearchResults, setAllSlashCommands], + [ + handleSayMessage, + handleAskMessage, + setFileSearchResults, + setAllSlashCommands, + setTokenUsage, + setRouterModels, + setApiConfiguration, + ], ) // Initialize extension host @@ -827,6 +874,8 @@ function AppInner({ cwd={workspacePath} reasoningEffort={reasoningEffort} version={version} + tokenUsage={tokenUsage} + contextWindow={contextWindow} /> diff --git a/apps/cli/src/ui/components/Header.tsx b/apps/cli/src/ui/components/Header.tsx index 9625193425..ab1458fa89 100644 --- a/apps/cli/src/ui/components/Header.tsx +++ b/apps/cli/src/ui/components/Header.tsx @@ -1,8 +1,11 @@ import { memo } from "react" import { Text, Box } from "ink" +import type { TokenUsage } from "@roo-code/types" + import { useTerminalSize } from "../hooks/TerminalSizeContext.js" import * as theme from "../utils/theme.js" +import MetricsDisplay from "./MetricsDisplay.js" interface HeaderProps { cwd: string @@ -10,6 +13,8 @@ interface HeaderProps { mode: string reasoningEffort?: string version: string + tokenUsage?: TokenUsage | null + contextWindow?: number } const ASCII_ROO = ` _,' ___ @@ -19,7 +24,7 @@ const ASCII_ROO = ` _,' ___ // \\\\ ,/' \`\\_,` -function Header({ model, cwd, mode, reasoningEffort, version }: HeaderProps) { +function Header({ model, cwd, mode, reasoningEffort, version, tokenUsage, contextWindow }: HeaderProps) { const { columns } = useTerminalSize() const homeDir = process.env.HOME || process.env.USERPROFILE || "" @@ -28,6 +33,9 @@ function Header({ model, cwd, mode, reasoningEffort, version }: HeaderProps) { const titlePart = `── ${title} ` const remainingDashes = Math.max(0, columns - titlePart.length) + // Only show metrics when we have token usage data + const showMetrics = tokenUsage && contextWindow && contextWindow > 0 + return ( @@ -43,6 +51,11 @@ function Header({ model, cwd, mode, reasoningEffort, version }: HeaderProps) { Mode: {mode} Model: {model} Reasoning: {reasoningEffort} + {showMetrics && ( + + + + )} diff --git a/apps/cli/src/ui/components/MetricsDisplay.tsx b/apps/cli/src/ui/components/MetricsDisplay.tsx new file mode 100644 index 0000000000..eed840b612 --- /dev/null +++ b/apps/cli/src/ui/components/MetricsDisplay.tsx @@ -0,0 +1,69 @@ +import { memo } from "react" +import { Text, Box } from "ink" + +import type { TokenUsage } from "@roo-code/types" + +import * as theme from "../utils/theme.js" +import ProgressBar from "./ProgressBar.js" + +interface MetricsDisplayProps { + tokenUsage: TokenUsage + contextWindow: number +} + +/** + * Formats a large number with K (thousands) or M (millions) suffix. + * + * Examples: + * - 1234 -> "1.2K" + * - 1234567 -> "1.2M" + * - 500 -> "500" + */ +function formatNumber(num: number): string { + if (num >= 1_000_000) { + return `${(num / 1_000_000).toFixed(1)}M` + } + if (num >= 1_000) { + return `${(num / 1_000).toFixed(1)}K` + } + return num.toString() +} + +/** + * Formats cost as currency with $ prefix. + * + * Examples: + * - 0.12345 -> "$0.12" + * - 1.5 -> "$1.50" + */ +function formatCost(cost: number): string { + return `$${cost.toFixed(2)}` +} + +/** + * Displays task metrics in a compact format: + * $0.12 │ ↓45.2K │ ↑8.7K │ Context: [████████░░░░] 62% + */ +function MetricsDisplay({ tokenUsage, contextWindow }: MetricsDisplayProps) { + const { totalCost, totalTokensIn, totalTokensOut, contextTokens } = tokenUsage + + return ( + + {formatCost(totalCost)} + • + + ↓ {formatNumber(totalTokensIn)} + + • + + ↑ {formatNumber(totalTokensOut)} + + • + Context: + + + ) +} + +export default memo(MetricsDisplay) +export { formatNumber, formatCost } diff --git a/apps/cli/src/ui/components/ProgressBar.tsx b/apps/cli/src/ui/components/ProgressBar.tsx new file mode 100644 index 0000000000..6c0e150dbe --- /dev/null +++ b/apps/cli/src/ui/components/ProgressBar.tsx @@ -0,0 +1,61 @@ +import { memo } from "react" +import { Text } from "ink" + +import * as theme from "../utils/theme.js" + +interface ProgressBarProps { + /** Current value (e.g., contextTokens) */ + value: number + /** Maximum value (e.g., contextWindow) */ + max: number + /** Width of the bar in characters (default: 16) */ + width?: number +} + +/** + * A progress bar component with color gradient based on fill percentage. + * + * Colors: + * - 0-50%: Green (safe zone) + * - 50-75%: Yellow (warning zone) + * - 75-100%: Red (danger zone) + * + * Visual example: [████████░░░░░░░░] 50% + */ +function ProgressBar({ value, max, width = 16 }: ProgressBarProps) { + // Calculate percentage, clamped to 0-100 + const percentage = max > 0 ? Math.min(100, Math.max(0, (value / max) * 100)) : 0 + + // Calculate how many blocks to fill + const filledBlocks = Math.round((percentage / 100) * width) + const emptyBlocks = width - filledBlocks + + // Determine color based on percentage + let barColor: string + if (percentage <= 50) { + barColor = theme.successColor // Green + } else if (percentage <= 75) { + barColor = theme.warningColor // Yellow + } else { + barColor = theme.errorColor // Red + } + + // Unicode block characters for smooth appearance + const filledChar = "█" + const emptyChar = "░" + + const filledPart = filledChar.repeat(filledBlocks) + const emptyPart = emptyChar.repeat(emptyBlocks) + + return ( + + [ + {filledPart} + + {emptyPart}] {Math.round(percentage)}% + + + ) +} + +export default memo(ProgressBar) diff --git a/apps/cli/src/ui/store.ts b/apps/cli/src/ui/store.ts index 8a67da4590..3111dde4b5 100644 --- a/apps/cli/src/ui/store.ts +++ b/apps/cli/src/ui/store.ts @@ -1,7 +1,15 @@ import { create } from "zustand" +import type { TokenUsage, ProviderSettings } from "@roo-code/types" + import type { TUIMessage, PendingAsk, FileSearchResult, SlashCommandResult } from "./types.js" +/** + * RouterModels type for context window lookup. + * Simplified version - we only need contextWindow from ModelInfo. + */ +export type RouterModels = Record> + /** * CLI application state. * @@ -23,6 +31,13 @@ interface CLIState { // Autocomplete data (from API/extension) fileSearchResults: FileSearchResult[] allSlashCommands: SlashCommandResult[] + + // Token usage metrics (from getApiMetrics) + tokenUsage: TokenUsage | null + + // Model info for context window lookup + routerModels: RouterModels | null + apiConfiguration: ProviderSettings | null } interface CLIActions { @@ -41,6 +56,11 @@ interface CLIActions { // Autocomplete data actions setFileSearchResults: (results: FileSearchResult[]) => void setAllSlashCommands: (commands: SlashCommandResult[]) => void + + // Metrics actions + setTokenUsage: (usage: TokenUsage | null) => void + setRouterModels: (models: RouterModels | null) => void + setApiConfiguration: (config: ProviderSettings | null) => void } const initialState: CLIState = { @@ -52,6 +72,9 @@ const initialState: CLIState = { error: null, fileSearchResults: [], allSlashCommands: [], + tokenUsage: null, + routerModels: null, + apiConfiguration: null, } export const useCLIStore = create((set) => ({ @@ -106,4 +129,7 @@ export const useCLIStore = create((set) => ({ reset: () => set(initialState), setFileSearchResults: (results) => set({ fileSearchResults: results }), setAllSlashCommands: (commands) => set({ allSlashCommands: commands }), + setTokenUsage: (usage) => set({ tokenUsage: usage }), + setRouterModels: (models) => set({ routerModels: models }), + setApiConfiguration: (config) => set({ apiConfiguration: config }), })) diff --git a/apps/cli/src/utils/getContextWindow.ts b/apps/cli/src/utils/getContextWindow.ts new file mode 100644 index 0000000000..d9915f18f7 --- /dev/null +++ b/apps/cli/src/utils/getContextWindow.ts @@ -0,0 +1,67 @@ +import type { ProviderSettings } from "@roo-code/types" + +import type { RouterModels } from "../ui/store.js" + +const DEFAULT_CONTEXT_WINDOW = 200_000 + +/** + * Looks up the context window size for the current model from routerModels. + * + * @param routerModels - The router models data containing model info per provider + * @param apiConfiguration - The current API configuration with provider and model ID + * @returns The context window size, or DEFAULT_CONTEXT_WINDOW (200K) if not found + */ +export function getContextWindow(routerModels: RouterModels | null, apiConfiguration: ProviderSettings | null): number { + if (!routerModels || !apiConfiguration) { + return DEFAULT_CONTEXT_WINDOW + } + + const provider = apiConfiguration.apiProvider + const modelId = getModelIdForProvider(apiConfiguration) + + if (!provider || !modelId) { + return DEFAULT_CONTEXT_WINDOW + } + + const providerModels = routerModels[provider] + const modelInfo = providerModels?.[modelId] + + return modelInfo?.contextWindow ?? DEFAULT_CONTEXT_WINDOW +} + +/** + * Gets the model ID from the API configuration based on the provider type. + * + * Different providers store their model ID in different fields of ProviderSettings. + */ +function getModelIdForProvider(config: ProviderSettings): string | undefined { + switch (config.apiProvider) { + case "openrouter": + return config.openRouterModelId + case "ollama": + return config.ollamaModelId + case "lmstudio": + return config.lmStudioModelId + case "openai": + return config.openAiModelId + case "requesty": + return config.requestyModelId + case "litellm": + return config.litellmModelId + case "deepinfra": + return config.deepInfraModelId + case "huggingface": + return config.huggingFaceModelId + case "unbound": + return config.unboundModelId + case "vercel-ai-gateway": + return config.vercelAiGatewayModelId + case "io-intelligence": + return config.ioIntelligenceModelId + default: + // For anthropic, bedrock, vertex, gemini, xai, groq, etc. + return config.apiModelId + } +} + +export { DEFAULT_CONTEXT_WINDOW } diff --git a/apps/cli/tsup.config.ts b/apps/cli/tsup.config.ts index d3bfb1407f..bc417d80cb 100644 --- a/apps/cli/tsup.config.ts +++ b/apps/cli/tsup.config.ts @@ -12,7 +12,7 @@ export default defineConfig({ js: "#!/usr/bin/env node", }, // Bundle workspace packages that export TypeScript - noExternal: ["@roo-code/types", "@roo-code/vscode-shim"], + noExternal: ["@roo-code/core", "@roo-code/core/message-utils", "@roo-code/types", "@roo-code/vscode-shim"], external: [ // Keep native modules external "@anthropic-ai/sdk", diff --git a/packages/core/package.json b/packages/core/package.json index 5151d88d95..41f2dfa95c 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -3,7 +3,10 @@ "description": "Platform agnostic core functionality for Roo Code.", "version": "0.0.0", "type": "module", - "exports": "./src/index.ts", + "exports": { + ".": "./src/index.ts", + "./message-utils": "./src/message-utils/index.ts" + }, "scripts": { "lint": "eslint src --ext=ts --max-warnings=0", "check-types": "tsc --noEmit", diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index fd7c93f68a..30c3ebe732 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -1 +1,2 @@ export * from "./custom-tools/index.js" +export * from "./message-utils/index.js" diff --git a/packages/core/src/message-utils/__tests__/consolidateApiRequests.spec.ts b/packages/core/src/message-utils/__tests__/consolidateApiRequests.spec.ts new file mode 100644 index 0000000000..363a0d0985 --- /dev/null +++ b/packages/core/src/message-utils/__tests__/consolidateApiRequests.spec.ts @@ -0,0 +1,122 @@ +// npx vitest run packages/core/src/message-utils/__tests__/consolidateApiRequests.spec.ts + +import type { ClineMessage } from "@roo-code/types" + +import { consolidateApiRequests } from "../consolidateApiRequests.js" + +describe("consolidateApiRequests", () => { + // Helper function to create a basic api_req_started message + const createApiReqStarted = (ts: number, data: Record = {}): ClineMessage => ({ + ts, + type: "say", + say: "api_req_started", + text: JSON.stringify(data), + }) + + // Helper function to create a basic api_req_finished message + const createApiReqFinished = (ts: number, data: Record = {}): ClineMessage => ({ + ts, + type: "say", + say: "api_req_finished", + text: JSON.stringify(data), + }) + + // Helper function to create a regular text message + const createTextMessage = (ts: number, text: string): ClineMessage => ({ + ts, + type: "say", + say: "text", + text, + }) + + it("should consolidate a matching pair of api_req_started and api_req_finished messages", () => { + const messages: ClineMessage[] = [ + createApiReqStarted(1000, { request: "GET /api/data" }), + createApiReqFinished(1001, { cost: 0.005 }), + ] + + const result = consolidateApiRequests(messages) + + expect(result.length).toBe(1) + expect(result[0].say).toBe("api_req_started") + + const parsedText = JSON.parse(result[0].text || "{}") + expect(parsedText.request).toBe("GET /api/data") + expect(parsedText.cost).toBe(0.005) + }) + + it("should handle messages with no api_req pairs", () => { + const messages: ClineMessage[] = [createTextMessage(1000, "Hello"), createTextMessage(1001, "World")] + + const result = consolidateApiRequests(messages) + + expect(result).toEqual(messages) + }) + + it("should handle empty messages array", () => { + const result = consolidateApiRequests([]) + expect(result).toEqual([]) + }) + + it("should handle single message array", () => { + const messages: ClineMessage[] = [createTextMessage(1000, "Hello")] + const result = consolidateApiRequests(messages) + expect(result).toEqual(messages) + }) + + it("should preserve non-api messages in the result", () => { + const messages: ClineMessage[] = [ + createTextMessage(1000, "Before"), + createApiReqStarted(1001, { request: "test" }), + createApiReqFinished(1002, { cost: 0.01 }), + createTextMessage(1003, "After"), + ] + + const result = consolidateApiRequests(messages) + + expect(result.length).toBe(3) + expect(result[0].text).toBe("Before") + expect(result[1].say).toBe("api_req_started") + expect(result[2].text).toBe("After") + }) + + it("should handle multiple api_req pairs", () => { + const messages: ClineMessage[] = [ + createApiReqStarted(1000, { request: "first" }), + createApiReqFinished(1001, { cost: 0.01 }), + createApiReqStarted(1002, { request: "second" }), + createApiReqFinished(1003, { cost: 0.02 }), + ] + + const result = consolidateApiRequests(messages) + + expect(result.length).toBe(2) + expect(JSON.parse(result[0].text || "{}").request).toBe("first") + expect(JSON.parse(result[1].text || "{}").request).toBe("second") + }) + + it("should handle orphan api_req_started without finish", () => { + const messages: ClineMessage[] = [ + createApiReqStarted(1000, { request: "orphan" }), + createTextMessage(1001, "Text"), + ] + + const result = consolidateApiRequests(messages) + + expect(result.length).toBe(2) + expect(result[0].say).toBe("api_req_started") + expect(JSON.parse(result[0].text || "{}").request).toBe("orphan") + }) + + it("should handle invalid JSON in message text", () => { + const messages: ClineMessage[] = [ + { ts: 1000, type: "say", say: "api_req_started", text: "invalid json" }, + createApiReqFinished(1001, { cost: 0.01 }), + ] + + const result = consolidateApiRequests(messages) + + // Should still consolidate, merging what it can + expect(result.length).toBe(1) + }) +}) diff --git a/packages/core/src/message-utils/__tests__/consolidateCommands.spec.ts b/packages/core/src/message-utils/__tests__/consolidateCommands.spec.ts new file mode 100644 index 0000000000..8b11a858de --- /dev/null +++ b/packages/core/src/message-utils/__tests__/consolidateCommands.spec.ts @@ -0,0 +1,145 @@ +// npx vitest run packages/core/src/message-utils/__tests__/consolidateCommands.spec.ts + +import type { ClineMessage } from "@roo-code/types" + +import { consolidateCommands, COMMAND_OUTPUT_STRING } from "../consolidateCommands.js" + +describe("consolidateCommands", () => { + describe("command sequences", () => { + it("should consolidate command and command_output messages", () => { + const messages: ClineMessage[] = [ + { type: "ask", ask: "command", text: "ls", ts: 1000 }, + { type: "ask", ask: "command_output", text: "file1.txt", ts: 1001 }, + { type: "ask", ask: "command_output", text: "file2.txt", ts: 1002 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(1) + expect(result[0].ask).toBe("command") + expect(result[0].text).toBe(`ls\n${COMMAND_OUTPUT_STRING}file1.txt\nfile2.txt`) + }) + + it("should handle multiple command sequences", () => { + const messages: ClineMessage[] = [ + { type: "ask", ask: "command", text: "ls", ts: 1000 }, + { type: "ask", ask: "command_output", text: "output1", ts: 1001 }, + { type: "ask", ask: "command", text: "pwd", ts: 1002 }, + { type: "ask", ask: "command_output", text: "output2", ts: 1003 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(2) + expect(result[0].text).toBe(`ls\n${COMMAND_OUTPUT_STRING}output1`) + expect(result[1].text).toBe(`pwd\n${COMMAND_OUTPUT_STRING}output2`) + }) + + it("should handle command without output", () => { + const messages: ClineMessage[] = [ + { type: "ask", ask: "command", text: "ls", ts: 1000 }, + { type: "say", say: "text", text: "some text", ts: 1001 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(2) + expect(result[0].ask).toBe("command") + expect(result[0].text).toBe("ls") + expect(result[1].say).toBe("text") + }) + + it("should handle duplicate outputs (ask and say with same text)", () => { + const messages: ClineMessage[] = [ + { type: "ask", ask: "command", text: "ls", ts: 1000 }, + { type: "ask", ask: "command_output", text: "same output", ts: 1001 }, + { type: "say", say: "command_output", text: "same output", ts: 1002 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(1) + expect(result[0].text).toBe(`ls\n${COMMAND_OUTPUT_STRING}same output`) + }) + }) + + describe("MCP server sequences", () => { + it("should consolidate use_mcp_server and mcp_server_response messages", () => { + const messages: ClineMessage[] = [ + { + type: "ask", + ask: "use_mcp_server", + text: JSON.stringify({ server: "test", tool: "myTool" }), + ts: 1000, + }, + { type: "say", say: "mcp_server_response", text: "response data", ts: 1001 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(1) + expect(result[0].ask).toBe("use_mcp_server") + const parsed = JSON.parse(result[0].text || "{}") + expect(parsed.server).toBe("test") + expect(parsed.response).toBe("response data") + }) + + it("should handle MCP request without response", () => { + const messages: ClineMessage[] = [ + { + type: "ask", + ask: "use_mcp_server", + text: JSON.stringify({ server: "test" }), + ts: 1000, + }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(1) + expect(result[0].ask).toBe("use_mcp_server") + }) + + it("should handle multiple MCP responses", () => { + const messages: ClineMessage[] = [ + { + type: "ask", + ask: "use_mcp_server", + text: JSON.stringify({ server: "test" }), + ts: 1000, + }, + { type: "say", say: "mcp_server_response", text: "response1", ts: 1001 }, + { type: "say", say: "mcp_server_response", text: "response2", ts: 1002 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(1) + const parsed = JSON.parse(result[0].text || "{}") + expect(parsed.response).toBe("response1\nresponse2") + }) + }) + + describe("mixed messages", () => { + it("should preserve non-command, non-MCP messages", () => { + const messages: ClineMessage[] = [ + { type: "say", say: "text", text: "before", ts: 1000 }, + { type: "ask", ask: "command", text: "ls", ts: 1001 }, + { type: "ask", ask: "command_output", text: "output", ts: 1002 }, + { type: "say", say: "text", text: "after", ts: 1003 }, + ] + + const result = consolidateCommands(messages) + + expect(result.length).toBe(3) + expect(result[0].text).toBe("before") + expect(result[1].text).toBe(`ls\n${COMMAND_OUTPUT_STRING}output`) + expect(result[2].text).toBe("after") + }) + + it("should handle empty array", () => { + const result = consolidateCommands([]) + expect(result).toEqual([]) + }) + }) +}) diff --git a/packages/core/src/message-utils/__tests__/consolidateTokenUsage.spec.ts b/packages/core/src/message-utils/__tests__/consolidateTokenUsage.spec.ts new file mode 100644 index 0000000000..e95ef61b07 --- /dev/null +++ b/packages/core/src/message-utils/__tests__/consolidateTokenUsage.spec.ts @@ -0,0 +1,246 @@ +// npx vitest run packages/core/src/message-utils/__tests__/consolidateTokenUsage.spec.ts + +import type { ClineMessage } from "@roo-code/types" + +import { consolidateTokenUsage, hasTokenUsageChanged, hasToolUsageChanged } from "../consolidateTokenUsage.js" + +describe("consolidateTokenUsage", () => { + // Helper function to create a basic api_req_started message + const createApiReqMessage = ( + ts: number, + data: { + tokensIn?: number + tokensOut?: number + cacheWrites?: number + cacheReads?: number + cost?: number + }, + ): ClineMessage => ({ + ts, + type: "say", + say: "api_req_started", + text: JSON.stringify(data), + }) + + describe("basic token accumulation", () => { + it("should accumulate tokens from a single message", () => { + const messages: ClineMessage[] = [createApiReqMessage(1000, { tokensIn: 100, tokensOut: 50, cost: 0.01 })] + + const result = consolidateTokenUsage(messages) + + expect(result.totalTokensIn).toBe(100) + expect(result.totalTokensOut).toBe(50) + expect(result.totalCost).toBe(0.01) + }) + + it("should accumulate tokens from multiple messages", () => { + const messages: ClineMessage[] = [ + createApiReqMessage(1000, { tokensIn: 100, tokensOut: 50, cost: 0.01 }), + createApiReqMessage(1001, { tokensIn: 200, tokensOut: 100, cost: 0.02 }), + ] + + const result = consolidateTokenUsage(messages) + + expect(result.totalTokensIn).toBe(300) + expect(result.totalTokensOut).toBe(150) + expect(result.totalCost).toBeCloseTo(0.03) + }) + + it("should handle cache writes and reads", () => { + const messages: ClineMessage[] = [ + createApiReqMessage(1000, { tokensIn: 100, tokensOut: 50, cacheWrites: 500, cacheReads: 200 }), + ] + + const result = consolidateTokenUsage(messages) + + expect(result.totalCacheWrites).toBe(500) + expect(result.totalCacheReads).toBe(200) + }) + + it("should handle empty messages array", () => { + const result = consolidateTokenUsage([]) + + expect(result.totalTokensIn).toBe(0) + expect(result.totalTokensOut).toBe(0) + expect(result.totalCost).toBe(0) + expect(result.contextTokens).toBe(0) + }) + }) + + describe("context tokens calculation", () => { + it("should calculate context tokens from the last API request", () => { + const messages: ClineMessage[] = [ + createApiReqMessage(1000, { tokensIn: 100, tokensOut: 50 }), + createApiReqMessage(1001, { tokensIn: 200, tokensOut: 100 }), + ] + + const result = consolidateTokenUsage(messages) + + // Context tokens = tokensIn + tokensOut from last message + expect(result.contextTokens).toBe(300) // 200 + 100 + }) + + it("should handle condense_context messages for context tokens", () => { + const messages: ClineMessage[] = [ + createApiReqMessage(1000, { tokensIn: 100, tokensOut: 50 }), + { + ts: 1001, + type: "say", + say: "condense_context", + contextCondense: { newContextTokens: 5000, cost: 0.05 }, + } as ClineMessage, + ] + + const result = consolidateTokenUsage(messages) + + expect(result.contextTokens).toBe(5000) + expect(result.totalCost).toBeCloseTo(0.05) + }) + }) + + describe("invalid data handling", () => { + it("should handle messages with invalid JSON", () => { + const messages: ClineMessage[] = [{ ts: 1000, type: "say", say: "api_req_started", text: "invalid json" }] + + // Should not throw + const result = consolidateTokenUsage(messages) + expect(result.totalTokensIn).toBe(0) + }) + + it("should skip non-api_req_started messages", () => { + const messages: ClineMessage[] = [ + { ts: 1000, type: "say", say: "text", text: "hello" }, + createApiReqMessage(1001, { tokensIn: 100, tokensOut: 50 }), + ] + + const result = consolidateTokenUsage(messages) + + expect(result.totalTokensIn).toBe(100) + expect(result.totalTokensOut).toBe(50) + }) + + it("should handle missing token values", () => { + const messages: ClineMessage[] = [createApiReqMessage(1000, { cost: 0.01 })] + + const result = consolidateTokenUsage(messages) + + expect(result.totalTokensIn).toBe(0) + expect(result.totalTokensOut).toBe(0) + expect(result.totalCost).toBe(0.01) + }) + }) +}) + +describe("hasTokenUsageChanged", () => { + it("should return true when snapshot is undefined", () => { + const current = { + totalTokensIn: 100, + totalTokensOut: 50, + totalCost: 0.01, + contextTokens: 150, + } + + expect(hasTokenUsageChanged(current, undefined)).toBe(true) + }) + + it("should return false when values are the same", () => { + const current = { + totalTokensIn: 100, + totalTokensOut: 50, + totalCost: 0.01, + contextTokens: 150, + } + const snapshot = { ...current } + + expect(hasTokenUsageChanged(current, snapshot)).toBe(false) + }) + + it("should return true when totalTokensIn changes", () => { + const current = { + totalTokensIn: 200, + totalTokensOut: 50, + totalCost: 0.01, + contextTokens: 150, + } + const snapshot = { + totalTokensIn: 100, + totalTokensOut: 50, + totalCost: 0.01, + contextTokens: 150, + } + + expect(hasTokenUsageChanged(current, snapshot)).toBe(true) + }) + + it("should return true when totalCost changes", () => { + const current = { + totalTokensIn: 100, + totalTokensOut: 50, + totalCost: 0.02, + contextTokens: 150, + } + const snapshot = { + totalTokensIn: 100, + totalTokensOut: 50, + totalCost: 0.01, + contextTokens: 150, + } + + expect(hasTokenUsageChanged(current, snapshot)).toBe(true) + }) +}) + +describe("hasToolUsageChanged", () => { + it("should return true when snapshot is undefined", () => { + const current = { + read_file: { attempts: 1, failures: 0 }, + } + + expect(hasToolUsageChanged(current, undefined)).toBe(true) + }) + + it("should return false when values are the same", () => { + const current = { + read_file: { attempts: 1, failures: 0 }, + } + const snapshot = { + read_file: { attempts: 1, failures: 0 }, + } + + expect(hasToolUsageChanged(current, snapshot)).toBe(false) + }) + + it("should return true when a tool is added", () => { + const current = { + read_file: { attempts: 1, failures: 0 }, + write_to_file: { attempts: 1, failures: 0 }, + } + const snapshot = { + read_file: { attempts: 1, failures: 0 }, + } + + expect(hasToolUsageChanged(current, snapshot)).toBe(true) + }) + + it("should return true when attempts change", () => { + const current = { + read_file: { attempts: 2, failures: 0 }, + } + const snapshot = { + read_file: { attempts: 1, failures: 0 }, + } + + expect(hasToolUsageChanged(current, snapshot)).toBe(true) + }) + + it("should return true when failures change", () => { + const current = { + read_file: { attempts: 1, failures: 1 }, + } + const snapshot = { + read_file: { attempts: 1, failures: 0 }, + } + + expect(hasToolUsageChanged(current, snapshot)).toBe(true) + }) +}) diff --git a/packages/core/src/message-utils/consolidateApiRequests.ts b/packages/core/src/message-utils/consolidateApiRequests.ts new file mode 100644 index 0000000000..ee538e015e --- /dev/null +++ b/packages/core/src/message-utils/consolidateApiRequests.ts @@ -0,0 +1,90 @@ +import type { ClineMessage } from "@roo-code/types" + +/** + * Consolidates API request start and finish messages in an array of ClineMessages. + * + * This function looks for pairs of 'api_req_started' and 'api_req_finished' messages. + * When it finds a pair, it consolidates them into a single message. + * The JSON data in the text fields of both messages are merged. + * + * @param messages - An array of ClineMessage objects to process. + * @returns A new array of ClineMessage objects with API requests consolidated. + * + * @example + * const messages = [ + * { type: "say", say: "api_req_started", text: '{"request":"GET /api/data"}', ts: 1000 }, + * { type: "say", say: "api_req_finished", text: '{"cost":0.005}', ts: 1001 } + * ]; + * const result = consolidateApiRequests(messages); + * // Result: [{ type: "say", say: "api_req_started", text: '{"request":"GET /api/data","cost":0.005}', ts: 1000 }] + */ +export function consolidateApiRequests(messages: ClineMessage[]): ClineMessage[] { + if (messages.length === 0) { + return [] + } + + if (messages.length === 1) { + return messages + } + + let isMergeNecessary = false + + for (const msg of messages) { + if (msg.type === "say" && (msg.say === "api_req_started" || msg.say === "api_req_finished")) { + isMergeNecessary = true + break + } + } + + if (!isMergeNecessary) { + return messages + } + + const result: ClineMessage[] = [] + const startedIndices: number[] = [] + + for (const message of messages) { + if (message.type !== "say" || (message.say !== "api_req_started" && message.say !== "api_req_finished")) { + result.push(message) + continue + } + + if (message.say === "api_req_started") { + // Add to result and track the index. + result.push(message) + startedIndices.push(result.length - 1) + continue + } + + // Find the most recent api_req_started that hasn't been consolidated. + const startIndex = startedIndices.length > 0 ? startedIndices.pop() : undefined + + if (startIndex !== undefined) { + const startMessage = result[startIndex] + if (!startMessage) continue + + let startData = {} + let finishData = {} + + try { + if (startMessage.text) { + startData = JSON.parse(startMessage.text) + } + } catch { + // Ignore JSON parse errors + } + + try { + if (message.text) { + finishData = JSON.parse(message.text) + } + } catch { + // Ignore JSON parse errors + } + + result[startIndex] = { ...startMessage, text: JSON.stringify({ ...startData, ...finishData }) } + } + } + + return result +} diff --git a/packages/core/src/message-utils/consolidateCommands.ts b/packages/core/src/message-utils/consolidateCommands.ts new file mode 100644 index 0000000000..32527d486a --- /dev/null +++ b/packages/core/src/message-utils/consolidateCommands.ts @@ -0,0 +1,160 @@ +import type { ClineMessage } from "@roo-code/types" + +import { safeJsonParse } from "./safeJsonParse.js" + +export const COMMAND_OUTPUT_STRING = "Output:" + +/** + * Consolidates sequences of command and command_output messages in an array of ClineMessages. + * Also consolidates sequences of use_mcp_server and mcp_server_response messages. + * + * This function processes an array of ClineMessages objects, looking for sequences + * where a 'command' message is followed by one or more 'command_output' messages, + * or where a 'use_mcp_server' message is followed by one or more 'mcp_server_response' messages. + * When such a sequence is found, it consolidates them into a single message, merging + * their text contents. + * + * @param messages - An array of ClineMessage objects to process. + * @returns A new array of ClineMessage objects with command and MCP sequences consolidated. + * + * @example + * const messages: ClineMessage[] = [ + * { type: 'ask', ask: 'command', text: 'ls', ts: 1625097600000 }, + * { type: 'ask', ask: 'command_output', text: 'file1.txt', ts: 1625097601000 }, + * { type: 'ask', ask: 'command_output', text: 'file2.txt', ts: 1625097602000 } + * ]; + * const result = consolidateCommands(messages); + * // Result: [{ type: 'ask', ask: 'command', text: 'ls\nfile1.txt\nfile2.txt', ts: 1625097600000 }] + */ +export function consolidateCommands(messages: ClineMessage[]): ClineMessage[] { + const consolidatedMessages = new Map() + const processedIndices = new Set() + + // Single pass through all messages + for (let i = 0; i < messages.length; i++) { + const msg = messages[i] + if (!msg) continue + + // Handle MCP server requests + if (msg.type === "ask" && msg.ask === "use_mcp_server") { + // Look ahead for MCP responses + const responses: string[] = [] + let j = i + 1 + + while (j < messages.length) { + const nextMsg = messages[j] + if (!nextMsg) { + j++ + continue + } + if (nextMsg.say === "mcp_server_response") { + responses.push(nextMsg.text || "") + processedIndices.add(j) + j++ + } else if (nextMsg.type === "ask" && nextMsg.ask === "use_mcp_server") { + // Stop if we encounter another MCP request + break + } else { + j++ + } + } + + if (responses.length > 0) { + // Parse the JSON from the message text + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const jsonObj = safeJsonParse(msg.text || "{}", {}) + + // Add the response to the JSON object + jsonObj.response = responses.join("\n") + + // Stringify the updated JSON object + const consolidatedText = JSON.stringify(jsonObj) + + consolidatedMessages.set(msg.ts, { ...msg, text: consolidatedText }) + } else { + // If there's no response, just keep the original message + consolidatedMessages.set(msg.ts, { ...msg }) + } + } + // Handle command sequences + else if (msg.type === "ask" && msg.ask === "command") { + let consolidatedText = msg.text || "" + let j = i + 1 + let previous: { type: "ask" | "say"; text: string } | undefined + let lastProcessedIndex = i + + while (j < messages.length) { + const currentMsg = messages[j] + if (!currentMsg) { + j++ + continue + } + const { type, ask, say, text = "" } = currentMsg + + if (type === "ask" && ask === "command") { + break // Stop if we encounter the next command. + } + + if (ask === "command_output" || say === "command_output") { + if (!previous) { + consolidatedText += `\n${COMMAND_OUTPUT_STRING}` + } + + const isDuplicate = previous && previous.type !== type && previous.text === text + + if (text.length > 0 && !isDuplicate) { + // Add a newline before adding the text if there's already content + if ( + previous && + consolidatedText.length > + consolidatedText.indexOf(COMMAND_OUTPUT_STRING) + COMMAND_OUTPUT_STRING.length + ) { + consolidatedText += "\n" + } + consolidatedText += text + } + + previous = { type, text } + processedIndices.add(j) + lastProcessedIndex = j + } + + j++ + } + + consolidatedMessages.set(msg.ts, { ...msg, text: consolidatedText }) + + // Only skip ahead if we actually processed command outputs + if (lastProcessedIndex > i) { + i = lastProcessedIndex + } + } + } + + // Build final result: filter out processed messages and use consolidated versions + const result: ClineMessage[] = [] + for (let i = 0; i < messages.length; i++) { + const msg = messages[i] + if (!msg) continue + + // Skip messages that were processed as outputs/responses + if (processedIndices.has(i)) { + continue + } + + // Skip command_output and mcp_server_response messages + if (msg.ask === "command_output" || msg.say === "command_output" || msg.say === "mcp_server_response") { + continue + } + + // Use consolidated version if available + const consolidatedMsg = consolidatedMessages.get(msg.ts) + if (consolidatedMsg) { + result.push(consolidatedMsg) + } else { + result.push(msg) + } + } + + return result +} diff --git a/packages/core/src/message-utils/consolidateTokenUsage.ts b/packages/core/src/message-utils/consolidateTokenUsage.ts new file mode 100644 index 0000000000..ca643afd17 --- /dev/null +++ b/packages/core/src/message-utils/consolidateTokenUsage.ts @@ -0,0 +1,157 @@ +import type { TokenUsage, ToolUsage, ToolName, ClineMessage } from "@roo-code/types" + +export type ParsedApiReqStartedTextType = { + tokensIn: number + tokensOut: number + cacheWrites: number + cacheReads: number + cost?: number // Only present if consolidateApiRequests has been called + apiProtocol?: "anthropic" | "openai" +} + +/** + * Consolidates token usage metrics from an array of ClineMessages. + * + * This function processes 'condense_context' messages and 'api_req_started' messages that have been + * consolidated with their corresponding 'api_req_finished' messages by the consolidateApiRequests function. + * It extracts and sums up the tokensIn, tokensOut, cacheWrites, cacheReads, and cost from these messages. + * + * @param messages - An array of ClineMessage objects to process. + * @returns A TokenUsage object containing totalTokensIn, totalTokensOut, totalCacheWrites, totalCacheReads, totalCost, and contextTokens. + * + * @example + * const messages = [ + * { type: "say", say: "api_req_started", text: '{"request":"GET /api/data","tokensIn":10,"tokensOut":20,"cost":0.005}', ts: 1000 } + * ]; + * const { totalTokensIn, totalTokensOut, totalCost } = consolidateTokenUsage(messages); + * // Result: { totalTokensIn: 10, totalTokensOut: 20, totalCost: 0.005 } + */ +export function consolidateTokenUsage(messages: ClineMessage[]): TokenUsage { + const result: TokenUsage = { + totalTokensIn: 0, + totalTokensOut: 0, + totalCacheWrites: undefined, + totalCacheReads: undefined, + totalCost: 0, + contextTokens: 0, + } + + // Calculate running totals. + messages.forEach((message) => { + if (message.type === "say" && message.say === "api_req_started" && message.text) { + try { + const parsedText: ParsedApiReqStartedTextType = JSON.parse(message.text) + const { tokensIn, tokensOut, cacheWrites, cacheReads, cost } = parsedText + + if (typeof tokensIn === "number") { + result.totalTokensIn += tokensIn + } + + if (typeof tokensOut === "number") { + result.totalTokensOut += tokensOut + } + + if (typeof cacheWrites === "number") { + result.totalCacheWrites = (result.totalCacheWrites ?? 0) + cacheWrites + } + + if (typeof cacheReads === "number") { + result.totalCacheReads = (result.totalCacheReads ?? 0) + cacheReads + } + + if (typeof cost === "number") { + result.totalCost += cost + } + } catch (error) { + console.error("Error parsing JSON:", error) + } + } else if (message.type === "say" && message.say === "condense_context") { + result.totalCost += message.contextCondense?.cost ?? 0 + } + }) + + // Calculate context tokens, from the last API request started or condense + // context message. + result.contextTokens = 0 + + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i] + if (!message) continue + + if (message.type === "say" && message.say === "api_req_started" && message.text) { + try { + const parsedText: ParsedApiReqStartedTextType = JSON.parse(message.text) + const { tokensIn, tokensOut } = parsedText + + // Since tokensIn now stores TOTAL input tokens (including cache tokens), + // we no longer need to add cacheWrites and cacheReads separately. + // This applies to both Anthropic and OpenAI protocols. + result.contextTokens = (tokensIn || 0) + (tokensOut || 0) + } catch { + // Ignore JSON parse errors + continue + } + } else if (message.type === "say" && message.say === "condense_context") { + result.contextTokens = message.contextCondense?.newContextTokens ?? 0 + } + if (result.contextTokens) { + break + } + } + + return result +} + +/** + * Check if token usage has changed by comparing relevant properties. + * @param current - Current token usage data + * @param snapshot - Previous snapshot to compare against + * @returns true if any relevant property has changed or snapshot is undefined + */ +export function hasTokenUsageChanged(current: TokenUsage, snapshot?: TokenUsage): boolean { + if (!snapshot) { + return true + } + + const keysToCompare: (keyof TokenUsage)[] = [ + "totalTokensIn", + "totalTokensOut", + "totalCacheWrites", + "totalCacheReads", + "totalCost", + "contextTokens", + ] + + return keysToCompare.some((key) => current[key] !== snapshot[key]) +} + +/** + * Check if tool usage has changed by comparing attempts and failures. + * @param current - Current tool usage data + * @param snapshot - Previous snapshot to compare against (undefined treated as empty) + * @returns true if any tool's attempts/failures have changed between current and snapshot + */ +export function hasToolUsageChanged(current: ToolUsage, snapshot?: ToolUsage): boolean { + // Treat undefined snapshot as empty object for consistent comparison + const effectiveSnapshot = snapshot ?? {} + + const currentKeys = Object.keys(current) as ToolName[] + const snapshotKeys = Object.keys(effectiveSnapshot) as ToolName[] + + // Check if number of tools changed + if (currentKeys.length !== snapshotKeys.length) { + return true + } + + // Check if any tool's stats changed + return currentKeys.some((key) => { + const currentTool = current[key] + const snapshotTool = effectiveSnapshot[key] + + if (!snapshotTool || !currentTool) { + return true + } + + return currentTool.attempts !== snapshotTool.attempts || currentTool.failures !== snapshotTool.failures + }) +} diff --git a/packages/core/src/message-utils/index.ts b/packages/core/src/message-utils/index.ts new file mode 100644 index 0000000000..b73600ea77 --- /dev/null +++ b/packages/core/src/message-utils/index.ts @@ -0,0 +1,12 @@ +export { + type ParsedApiReqStartedTextType, + consolidateTokenUsage, + hasTokenUsageChanged, + hasToolUsageChanged, +} from "./consolidateTokenUsage.js" + +export { consolidateApiRequests } from "./consolidateApiRequests.js" + +export { consolidateCommands, COMMAND_OUTPUT_STRING } from "./consolidateCommands.js" + +export { safeJsonParse } from "./safeJsonParse.js" diff --git a/packages/core/src/message-utils/safeJsonParse.ts b/packages/core/src/message-utils/safeJsonParse.ts new file mode 100644 index 0000000000..c60f8b3b84 --- /dev/null +++ b/packages/core/src/message-utils/safeJsonParse.ts @@ -0,0 +1,20 @@ +/** + * Safely parses JSON without crashing on invalid input. + * + * @param jsonString The string to parse + * @param defaultValue Value to return if parsing fails + * @returns Parsed JSON object or defaultValue if parsing fails + */ +export function safeJsonParse(jsonString: string | null | undefined, defaultValue?: T): T | undefined { + if (!jsonString) { + return defaultValue + } + + try { + return JSON.parse(jsonString) as T + } catch (error) { + // Log the error to the console for debugging. + console.error("Error parsing JSON:", error) + return defaultValue + } +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 4aa5aae148..012442ac54 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -85,6 +85,9 @@ importers: '@inkjs/ui': specifier: ^2.0.0 version: 2.0.0(ink@6.6.0(@types/react@18.3.23)(react@19.2.3)) + '@roo-code/core': + specifier: workspace:^ + version: link:../../packages/core '@roo-code/types': specifier: workspace:^ version: link:../../packages/types