move token-counter to agent-runtime

Author: charleslien

backend/src/llm-apis/relace-api.ts

@@ -1,10 +1,10 @@
+import { countTokens } from '@codebuff/agent-runtime/util/token-counter'
 import { models } from '@codebuff/common/old-constants'
 import { buildArray } from '@codebuff/common/util/array'
 import { parseMarkdownCodeBlock } from '@codebuff/common/util/file'
 import { env } from '@codebuff/internal'
 
 import { saveMessage } from '../llm-apis/message-cost-tracker'
-import { countTokens } from '../util/token-counter'
 
 import type { PromptAiSdkFn } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'

backend/src/process-file-block.ts

@@ -1,3 +1,4 @@
+import { countTokens } from '@codebuff/agent-runtime/util/token-counter'
 import { models } from '@codebuff/common/old-constants'
 import { cleanMarkdownCodeBlock } from '@codebuff/common/util/file'
 import { hasLazyEdit } from '@codebuff/common/util/string'
@@ -8,7 +9,6 @@ import {
   parseAndGetDiffBlocksSingleFile,
   retryDiffBlocksPrompt,
 } from './generate-diffs-prompt'
-import { countTokens } from './util/token-counter'
 
 import type { PromptAiSdkFn } from '@codebuff/common/types/contracts/llm'
 import type { Logger } from '@codebuff/common/types/contracts/logger'

backend/src/run-agent-step.ts

@@ -1,3 +1,4 @@
+import { countTokensJson } from '@codebuff/agent-runtime/util/token-counter'
 import { insertTrace } from '@codebuff/bigquery'
 import { trackEvent } from '@codebuff/common/analytics'
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
@@ -23,17 +24,16 @@ import {
   messagesWithSystem,
   expireMessages,
 } from './util/messages'
-import { countTokensJson } from './util/token-counter'
 import { getRequestContext } from './websockets/request-context'
 
 import type { AgentResponseTrace } from '@codebuff/bigquery'
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
+import type { SendActionFn } from '@codebuff/common/types/contracts/client'
 import type {
   AddAgentStepFn,
   FinishAgentRunFn,
   StartAgentRunFn,
 } from '@codebuff/common/types/contracts/database'
-import type { SendActionFn } from '@codebuff/common/types/contracts/client'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { Message } from '@codebuff/common/types/messages/codebuff-message'

backend/src/system-prompt/search-system-prompt.ts

@@ -1,3 +1,7 @@
+import {
+  countTokens,
+  countTokensJson,
+} from '@codebuff/agent-runtime/util/token-counter'
 import { insertTrace } from '@codebuff/bigquery'
 import { buildArray } from '@codebuff/common/util/array'
 
@@ -6,7 +10,6 @@ import {
   getProjectFileTreePrompt,
   getSystemInfoPrompt,
 } from './prompts'
-import { countTokens, countTokensJson } from '../util/token-counter'
 
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ProjectFileContext } from '@codebuff/common/util/file'

backend/src/system-prompt/truncate-file-tree.ts

@@ -1,13 +1,14 @@
+import {
+  countTokens,
+  countTokensJson,
+} from '@codebuff/agent-runtime/util/token-counter'
 import {
   printFileTree,
   printFileTreeWithTokens,
 } from '@codebuff/common/util/file'
 import { sampleSizeWithSeed } from '@codebuff/common/util/random'
 
-import { countTokens, countTokensJson } from '../util/token-counter'
-
 import type { Logger } from '@codebuff/common/types/contracts/logger'
-
 import type {
   FileTreeNode,
   ProjectFileContext,

backend/src/tools/handlers/tool/find-files.ts

@@ -1,3 +1,7 @@
+import {
+  countTokens,
+  countTokensJson,
+} from '@codebuff/agent-runtime/util/token-counter'
 import { insertTrace } from '@codebuff/bigquery'
 
 import {
@@ -7,7 +11,6 @@ import {
 import { getFileReadingUpdates } from '../../../get-file-reading-updates'
 import { getSearchSystemPrompt } from '../../../system-prompt/search-system-prompt'
 import { renderReadFilesResult } from '../../../util/parse-tool-call-xml'
-import { countTokens, countTokensJson } from '../../../util/token-counter'
 
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
 import type { GetExpandedFileContextForTrainingBlobTrace } from '@codebuff/bigquery'

backend/src/util/__tests__/messages.test.ts

@@ -1,3 +1,4 @@
+import * as tokenCounter from '@codebuff/agent-runtime/util/token-counter'
 import {
   afterEach,
   beforeEach,
@@ -13,7 +14,6 @@ import {
   messagesWithSystem,
   getPreviouslyReadFiles,
 } from '../messages'
-import * as tokenCounter from '../token-counter'
 
 import type { CodebuffToolMessage } from '@codebuff/common/tools/list'
 import type {
@@ -292,7 +292,12 @@ describe('trimMessagesToFitTokenLimit', () => {
   it('handles empty messages array', () => {
     const maxTotalTokens = 200
     const systemTokens = 100
-    const result = trimMessagesToFitTokenLimit({ messages: [], systemTokens, maxTotalTokens, logger })
+    const result = trimMessagesToFitTokenLimit({
+      messages: [],
+      systemTokens,
+      maxTotalTokens,
+      logger,
+    })
 
     expect(result).toEqual([])
   })
@@ -315,7 +320,12 @@ describe('trimMessagesToFitTokenLimit', () => {
         },
       ]
 
-      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 1000, logger })
+      const result = trimMessagesToFitTokenLimit({
+        messages,
+        systemTokens: 0,
+        maxTotalTokens: 1000,
+        logger,
+      })
 
       // Should contain the kept messages
       const keptMessages = result.filter(
@@ -345,7 +355,12 @@ describe('trimMessagesToFitTokenLimit', () => {
         },
       ] as Message[]
 
-      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 10000, logger })
+      const result = trimMessagesToFitTokenLimit({
+        messages,
+        systemTokens: 0,
+        maxTotalTokens: 10000,
+        logger,
+      })
 
       // Should be unchanged when under token limit
       expect(result).toHaveLength(2)
@@ -361,7 +376,12 @@ describe('trimMessagesToFitTokenLimit', () => {
         { role: 'user', content: 'Keep this', keepDuringTruncation: true },
       ]
 
-      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 1000, logger })
+      const result = trimMessagesToFitTokenLimit({
+        messages,
+        systemTokens: 0,
+        maxTotalTokens: 1000,
+        logger,
+      })
 
       // Should only have one replacement message for consecutive removals
       const replacementMessages = result.filter(
@@ -391,7 +411,12 @@ describe('trimMessagesToFitTokenLimit', () => {
         { role: 'user', content: 'C'.repeat(100) }, // Might be kept
       ]
 
-      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 2000, logger })
+      const result = trimMessagesToFitTokenLimit({
+        messages,
+        systemTokens: 0,
+        maxTotalTokens: 2000,
+        logger,
+      })
 
       // Should preserve the keepDuringTruncation message
       const keptMessage = result.find(
@@ -415,7 +440,12 @@ describe('trimMessagesToFitTokenLimit', () => {
         { role: 'user', content: 'C'.repeat(800) }, // Large message to force truncation
       ]
 
-      const result = trimMessagesToFitTokenLimit({ messages, systemTokens: 0, maxTotalTokens: 500, logger })
+      const result = trimMessagesToFitTokenLimit({
+        messages,
+        systemTokens: 0,
+        maxTotalTokens: 500,
+        logger,
+      })
 
       // Should keep both marked messages
       const keptMessages = result.filter(

backend/src/util/messages.ts

@@ -1,18 +1,19 @@
 import { AssertionError } from 'assert'
 
+import { countTokensJson } from '@codebuff/agent-runtime/util/token-counter'
 import { buildArray } from '@codebuff/common/util/array'
 import { getErrorObject } from '@codebuff/common/util/error'
 import { closeXml } from '@codebuff/common/util/xml'
 import { cloneDeep, isEqual } from 'lodash'
 
 import { simplifyTerminalCommandResults } from './simplify-tool-results'
-import { countTokensJson } from './token-counter'
 
 import type { System } from '../llm-apis/claude'
 import type {
   CodebuffToolMessage,
   CodebuffToolOutput,
 } from '@codebuff/common/tools/list'
+import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type {
   Message,
   ToolMessage,
@@ -21,7 +22,6 @@ import type {
   TextPart,
   ImagePart,
 } from '@codebuff/common/types/messages/content-part'
-import type { Logger } from '@codebuff/common/types/contracts/logger'
 
 export function messagesWithSystem(params: {
   messages: Message[]