sdk: implement run, allow continuing from previous run

Author: jahooma

sdk/src/client.ts

@@ -5,8 +5,14 @@ import { CODEBUFF_BINARY } from './constants'
 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 {
+  PromptResponseSchema,
+  type ServerAction,
+} from '../../common/src/actions'
+import {
+  getInitialSessionState,
+  SessionState,
+} from '../../common/src/types/session-state'
 import { getFiles } from '../../npm-app/src/project-files'
 
 export type ClientToolName =
@@ -30,11 +36,20 @@ export type CodebuffClientOptions = {
   }
 }
 
+type RunState = {
+  sessionState: SessionState
+  toolResults: Extract<ServerAction, { type: 'prompt-response' }>['toolResults']
+}
+
 export class CodebuffClient {
+  public cwd: string
   private readonly websocketHandler: WebSocketHandler
   private readonly overrideTools: CodebuffClientOptions['overrideTools']
   private readonly fingerprintId = `codebuff-sdk-${Math.random().toString(36).substring(2, 15)}`
-  public cwd: string
+  private readonly promptIdToResolveResponse: Record<
+    string,
+    { resolve: (response: any) => void; reject: (error: any) => void }
+  > = {}
 
   constructor({ cwd, onError, overrideTools }: CodebuffClientOptions) {
     // TODO: download binary automatically
@@ -68,45 +83,103 @@ export class CodebuffClient {
       onResponseChunk: async () => {},
       onSubagentResponseChunk: async () => {},
 
-      onPromptResponse: async () => {},
+      onPromptResponse: this.handlePromptResponse.bind(this),
     })
   }
 
-  public async runNewChat({
+  /**
+   * Run an agent.
+   *
+   * Pass an agent id, a prompt, and an event handler, plus options.
+   *
+   * Returns the state of the run, which can be passed to a subsequent run to continue the run.
+   *
+   * @param agent - The agent to run, e.g. 'base' or 'codebuff/file-picker@0.0.1'
+   * @param prompt - The user prompt, e.g. 'Add a console.log to the index file'
+   * @param params - (Optional) The parameters to pass to the agent.
+   * @param handleEvent - A function to handle events.
+   *
+   * @param previousState - (Optional) Continue a previous run with the return value of a previous run.
+   *
+   * @param allFiles - (Optional) All the files in the project, in an object of file path to file content. Improves codebuff's ability to locate files.
+   * @param knowledgeFiles - (Optional) The knowledge files to pass to the agent.
+   * @param agentTemplates - (Optional) The agent templates to pass to the agent.
+   * @param maxAgentSteps - (Optional) The maximum number of agent steps the main agent can run before stopping.
+   */
+  public async run({
     agent,
     prompt,
     params,
     handleEvent,
+    previousState,
     allFiles,
     knowledgeFiles,
-    agentTemplates,
+    agentConfig,
+    maxAgentSteps,
   }: {
     agent: string
     prompt: string
     params?: Record<string, any>
     handleEvent: (event: any) => void
+    previousState?: RunState
     allFiles?: Record<string, string>
     knowledgeFiles?: Record<string, string>
-    agentTemplates?: Record<string, any>
-  }): Promise<{
-    agentId: string
-  }> {
+    agentConfig?: Record<string, any>
+    maxAgentSteps?: number
+  }): Promise<RunState> {
+    const promptId = Math.random().toString(36).substring(2, 15)
+    const sessionState =
+      previousState?.sessionState ??
+      initialSessionState(this.cwd, {
+        knowledgeFiles,
+        agentConfig,
+        allFiles,
+        maxAgentSteps,
+      })
+    const toolResults = previousState?.toolResults ?? []
     this.websocketHandler.sendInput({
-      promptId: Math.random().toString(36).substring(2, 15),
+      promptId,
       prompt,
       promptParams: params,
       fingerprintId: this.fingerprintId,
       costMode: 'normal',
-      sessionState: initialSessionState(this.cwd, {
-        knowledgeFiles,
-        agentTemplates,
-        allFiles,
-      }),
-      toolResults: [],
+      sessionState,
+      toolResults,
       agentId: agent,
     })
 
-    return new Promise((resolve) => {})
+    return new Promise<RunState>((resolve, reject) => {
+      this.promptIdToResolveResponse[promptId] = { resolve, reject }
+    })
+  }
+
+  private async handlePromptResponse(
+    action: Extract<ServerAction, { type: 'prompt-response' }>,
+  ) {
+    const promiseActions =
+      this.promptIdToResolveResponse[action?.promptId ?? '']
+
+    const parsedAction = PromptResponseSchema.safeParse(action)
+    if (!parsedAction.success) {
+      const message = [
+        'Received invalid prompt response from server:',
+        JSON.stringify(parsedAction.error.errors),
+        'If this issues persists, please contact support@codebuff.com',
+      ].join('\n')
+      if (promiseActions) {
+        promiseActions.reject(new Error(message))
+      }
+      return
+    }
+
+    if (promiseActions) {
+      const { sessionState, toolResults } = action
+      const state: RunState = {
+        sessionState,
+        toolResults,
+      }
+      promiseActions.resolve(state)
+    }
   }
 
   private async readFiles(filePath: string[]) {
@@ -164,25 +237,27 @@ export class CodebuffClient {
     }
   }
 }
+
 function initialSessionState(
   cwd: string,
   options: {
     allFiles?: Record<string, string>
     knowledgeFiles?: Record<string, string>
-    agentTemplates?: Record<string, any>
+    agentConfig?: Record<string, any>
+    maxAgentSteps?: number
   },
 ) {
-  const { knowledgeFiles = {}, agentTemplates = {} } = options
+  const { knowledgeFiles = {}, agentConfig = {} } = options
 
-  return getInitialSessionState({
+  const initialState = getInitialSessionState({
     projectRoot: cwd,
     cwd,
     fileTree: [],
     fileTokenScores: {},
     tokenCallers: {},
     knowledgeFiles,
     userKnowledgeFiles: {},
-    agentTemplates,
+    agentTemplates: agentConfig,
     gitChanges: {
       status: '',
       diff: '',
@@ -200,4 +275,10 @@ function initialSessionState(
       cpus: 16,
     },
   })
+
+  if (options.maxAgentSteps) {
+    initialState.mainAgentState.stepsRemaining = options.maxAgentSteps
+  }
+
+  return initialState
 }