Tool tweaks (#380)

Author: charleslien

backend/src/__tests__/main-prompt.integration.test.ts

@@ -7,7 +7,7 @@ import { getToolCallString } from '@codebuff/common/tools/utils'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
   assistantMessage,
-  toolJsonContent,
+  jsonToolResult,
 } from '@codebuff/common/util/messages'
 import { afterEach, beforeEach, describe, expect, it, mock } from 'bun:test'
 
@@ -342,12 +342,10 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
         role: 'tool',
         toolName: 'read_files',
         toolCallId: 'test-id',
-        content: [
-          toolJsonContent({
-            path: 'src/util/messages.ts',
-            content: initialContent,
-          }),
-        ],
+        content: jsonToolResult({
+          path: 'src/util/messages.ts',
+          content: initialContent,
+        }),
       },
     )
 
@@ -415,12 +413,10 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
           role: 'tool',
           toolName: 'read_files',
           toolCallId: 'test-id',
-          content: [
-            toolJsonContent({
-              path: 'packages/backend/src/index.ts',
-              content: initialContent,
-            }),
-          ],
+          content: jsonToolResult({
+            path: 'packages/backend/src/index.ts',
+            content: initialContent,
+          }),
         },
       )
 

backend/src/websockets/server.ts

@@ -1,5 +1,4 @@
 import { setSessionConnected } from '@codebuff/agent-runtime/live-user-inputs'
-import { CLIENT_MESSAGE_SCHEMA } from '@codebuff/common/websockets/websocket-schema'
 import { isError } from 'lodash'
 import { WebSocketServer } from 'ws'
 
@@ -50,7 +49,7 @@ async function processMessage(params: {
   }
 
   try {
-    const msg = CLIENT_MESSAGE_SCHEMA.parse(messageObj)
+    const msg = messageObj
     const { type, txid } = msg
     switch (type) {
       case 'subscribe': {

backend/src/websockets/websocket-action.ts

@@ -2,12 +2,13 @@ import {
   cancelUserInput,
   startUserInput,
 } from '@codebuff/agent-runtime/live-user-inputs'
+import { callMainPrompt } from '@codebuff/agent-runtime/main-prompt'
 import { calculateUsageAndBalance } from '@codebuff/billing'
 import { trackEvent } from '@codebuff/common/analytics'
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
+import { getErrorObject } from '@codebuff/common/util/error'
 import db from '@codebuff/internal/db/index'
 import * as schema from '@codebuff/internal/db/schema'
-import { getErrorObject } from '@codebuff/common/util/error'
 import { eq } from 'drizzle-orm'
 
 import { protec } from './middleware'
@@ -22,7 +23,6 @@ import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
 import type { ClientMessage } from '@codebuff/common/websockets/websocket-schema'
 import type { WebSocket } from 'ws'
-import { callMainPrompt } from '@codebuff/agent-runtime/main-prompt'
 
 /**
  * Generates a usage response object for the client
@@ -330,9 +330,12 @@ export const onWebsocketAction = async (params: {
 }
 
 // Register action handlers
-subscribeToAction('prompt', protec.run({ baseAction: onPrompt }))
-subscribeToAction('init', protec.run({ baseAction: onInit, silent: true }))
+subscribeToAction('prompt', protec.run<'prompt'>({ baseAction: onPrompt }))
+subscribeToAction(
+  'init',
+  protec.run<'init'>({ baseAction: onInit, silent: true }),
+)
 subscribeToAction(
   'cancel-user-input',
-  protec.run({ baseAction: onCancelUserInput }),
+  protec.run<'cancel-user-input'>({ baseAction: onCancelUserInput }),
 )