sdk, in progress client: runNewChat, handle tools

Author: jahooma

sdk/src/client.ts

@@ -1,60 +1,203 @@
 import { execFileSync } from 'child_process'
+import os from 'os'
 
 import { CODEBUFF_BINARY } from './constants'
-import { processStream } from './process-stream'
+import { changeFile } from './tools/change-file'
+import { WebSocketHandler } from './websocket-client'
 import { API_KEY_ENV_VAR } from '../../common/src/constants'
+import type { ServerAction } from '../../common/src/actions'
+import { getInitialSessionState } from '../../common/src/types/session-state'
+import { getFiles } from '../../npm-app/src/project-files'
+
+export type ClientToolName =
+  | 'read_files'
+  | 'write_file'
+  | 'str_replace'
+  | 'run_terminal_command'
+
+export type CodebuffClientOptions = {
+  cwd: string
+  onError: (error: { message: string }) => void
+  overrideTools: Record<
+    ClientToolName,
+    (
+      args: Extract<ServerAction, { type: 'tool-call-request' }>['args'],
+    ) => Promise<{ toolResultMessage: string }>
+  > & {
+    readFiles: (
+      filePath: string[],
+    ) => Promise<{ files: Record<string, string | null> }>
+  }
+}
 
-/** @deprecated Migrate to WebSocketHandler */
 export class CodebuffClient {
+  private readonly websocketHandler: WebSocketHandler
+  private readonly overrideTools: CodebuffClientOptions['overrideTools']
+  private readonly fingerprintId = `codebuff-sdk-${Math.random().toString(36).substring(2, 15)}`
   public cwd: string
 
-  constructor({ cwd }: { cwd: string }) {
+  constructor({ cwd, onError, overrideTools }: CodebuffClientOptions) {
     // TODO: download binary automatically
     if (execFileSync('which', [CODEBUFF_BINARY]).toString().trim() === '') {
       throw new Error(
-        'Codebuff binary not found. Please run "npm i -g codebuff"',
+        `Could not find ${CODEBUFF_BINARY} in PATH. Please run "npm i -g codebuff" to install the codebuff.`,
       )
     }
     if (!process.env[API_KEY_ENV_VAR]) {
       throw new Error(
         `Codebuff API key not found. Please set the ${API_KEY_ENV_VAR} environment variable.`,
       )
     }
+    const apiKey = process.env[API_KEY_ENV_VAR]
 
     this.cwd = cwd
+    this.overrideTools = overrideTools
+    this.websocketHandler = new WebSocketHandler({
+      apiKey,
+      onWebsocketError: () => {},
+      onWebsocketReconnect: () => {},
+      onRequestReconnect: async () => {},
+      onResponseError: async (error) => {
+        onError({ message: error.message })
+      },
+      readFiles: this.readFiles.bind(this),
+      handleToolCall: this.handleToolCall.bind(this),
+      onCostResponse: async () => {},
+      onUsageResponse: async () => {},
+
+      onResponseChunk: async () => {},
+      onSubagentResponseChunk: async () => {},
+
+      onPromptResponse: async () => {},
+    })
   }
 
   public async runNewChat({
     agent,
     prompt,
     params,
     handleEvent,
+    allFiles,
+    knowledgeFiles,
+    agentTemplates,
   }: {
     agent: string
     prompt: string
     params?: Record<string, any>
     handleEvent: (event: any) => void
+    allFiles?: Record<string, string>
+    knowledgeFiles?: Record<string, string>
+    agentTemplates?: Record<string, any>
   }): Promise<{
     agentId: string
   }> {
-    const args = [prompt, '-p', '--agent', agent]
-    if (prompt) {
-      args.push(prompt)
-    }
-    if (params) {
-      args.push('--params', JSON.stringify(params))
-    }
-    if (this.cwd) {
-      args.push('--cwd', this.cwd)
-    }
-
-    await processStream({
-      codebuffArgs: args,
-      handleEvent,
+    this.websocketHandler.sendInput({
+      promptId: Math.random().toString(36).substring(2, 15),
+      prompt,
+      promptParams: params,
+      fingerprintId: this.fingerprintId,
+      costMode: 'normal',
+      sessionState: initialSessionState(this.cwd, {
+        knowledgeFiles,
+        agentTemplates,
+        allFiles,
+      }),
+      toolResults: [],
+      agentId: agent,
     })
 
+    return new Promise((resolve) => {})
+  }
+
+  private async readFiles(filePath: string[]) {
+    const override = this.overrideTools.readFiles
+    if (override) {
+      const overrideResult = await override(filePath)
+      return overrideResult.files
+    }
+    return getFiles(filePath)
+  }
+
+  private async handleToolCall(
+    action: Extract<ServerAction, { type: 'tool-call-request' }>,
+  ) {
+    const toolName = action.toolName
+    const args = action.args
+    let result: string
+    try {
+      const override = this.overrideTools[toolName as ClientToolName]
+      if (override) {
+        const overrideResult = await override(args)
+        result = overrideResult.toolResultMessage
+      } else if (toolName === 'end_turn') {
+        result = ''
+      } else if (toolName === 'write_file' || toolName === 'str_replace') {
+        const r = changeFile(args, this.cwd)
+        result = r.toolResultMessage
+      } else if (toolName === 'run_terminal_command') {
+        throw new Error(
+          'run_terminal_command not implemented in SDK yet; please provide an override.',
+        )
+      } else {
+        throw new Error(
+          `Tool not implemented in sdk. Please provide an override or modify your agent to not use this tool: ${toolName}`,
+        )
+      }
+    } catch (error) {
+      return {
+        type: 'tool-call-response',
+        requestId: action.requestId,
+        success: false,
+        result:
+          error && typeof error === 'object' && 'message' in error
+            ? error.message
+            : typeof error === 'string'
+              ? error
+              : 'Unknown error',
+      }
+    }
     return {
-      agentId: agent,
+      type: 'tool-call-response',
+      requestId: action.requestId,
+      success: true,
+      result,
     }
   }
 }
+function initialSessionState(
+  cwd: string,
+  options: {
+    allFiles?: Record<string, string>
+    knowledgeFiles?: Record<string, string>
+    agentTemplates?: Record<string, any>
+  },
+) {
+  const { knowledgeFiles = {}, agentTemplates = {} } = options
+
+  return getInitialSessionState({
+    projectRoot: cwd,
+    cwd,
+    fileTree: [],
+    fileTokenScores: {},
+    tokenCallers: {},
+    knowledgeFiles,
+    userKnowledgeFiles: {},
+    agentTemplates,
+    gitChanges: {
+      status: '',
+      diff: '',
+      diffCached: '',
+      lastCommitMessages: '',
+    },
+    changesSinceLastChat: {},
+    shellConfigFiles: {},
+    systemInfo: {
+      platform: process.platform,
+      shell: 'bash',
+      nodeVersion: process.version,
+      arch: process.arch,
+      homedir: os.homedir(),
+      cpus: 16,
+    },
+  })
+}

sdk/src/tools/change-file.ts

@@ -0,0 +1,86 @@
+import z from 'zod'
+import fs from 'fs'
+import path from 'path'
+import { applyPatch } from '@codebuff/common/util/patch'
+
+const FileChangeSchema = z.object({
+  type: z.enum(['patch', 'file']),
+  path: z.string(),
+  content: z.string(),
+})
+
+export function changeFile(
+  parameters: unknown,
+  cwd: string,
+): { toolResultMessage: string } {
+  const fileChange = FileChangeSchema.parse(parameters)
+  const lines = fileChange.content.split('\n')
+
+  const { created, modified, invalid } = applyChanges(cwd, [fileChange])
+
+  const results: string[] = []
+
+  for (const file of created) {
+    results.push(
+      `Created ${file} successfully. Changes made:\n${lines.join('\n')}`,
+    )
+  }
+
+  for (const file of modified) {
+    results.push(
+      `Wrote to ${file} successfully. Changes made:\n${lines.join('\n')}`,
+    )
+  }
+
+  for (const file of invalid) {
+    results.push(
+      `Failed to write to ${file}; file path caused an error or file could not be written`,
+    )
+  }
+
+  return { toolResultMessage: results.join('\n') }
+}
+
+function applyChanges(
+  projectRoot: string,
+  changes: {
+    type: 'patch' | 'file'
+    path: string
+    content: string
+  }[],
+) {
+  const created: string[] = []
+  const modified: string[] = []
+  const invalid: string[] = []
+
+  for (const change of changes) {
+    const { path: filePath, content, type } = change
+    try {
+      const fullPath = path.join(projectRoot, filePath)
+      const fileExists = fs.existsSync(fullPath)
+      if (!fileExists) {
+        const dirPath = path.dirname(fullPath)
+        fs.mkdirSync(dirPath, { recursive: true })
+      }
+
+      if (type === 'file') {
+        fs.writeFileSync(fullPath, content)
+      } else {
+        const oldContent = fs.readFileSync(fullPath, 'utf-8')
+        const newContent = applyPatch(oldContent, content)
+        fs.writeFileSync(fullPath, newContent)
+      }
+
+      if (fileExists) {
+        modified.push(filePath)
+      } else {
+        created.push(filePath)
+      }
+    } catch (error) {
+      console.error(`Failed to apply patch to ${filePath}:`, error, content)
+      invalid.push(filePath)
+    }
+  }
+
+  return { created, modified, invalid }
+}

sdk/src/websocket-client.ts

@@ -157,7 +157,6 @@ export class WebSocketHandler {
     return {
       ...({
         type: 'prompt',
-        fingerprintId: 'codebuff-sdk',
       } as const),
       authToken: this.apiKey,
     }