SDK (#231)

charleslien · web-flow · commit 43a1cd4e2845 · 2025-08-05T14:11:34.000-07:00
diff --git a/backend/src/__tests__/main-prompt.integration.test.ts b/backend/src/__tests__/main-prompt.integration.test.ts
@@ -21,7 +21,7 @@ import { logger } from '../util/logger'
 import { renderReadFilesResult } from '../util/parse-tool-call-xml'
 import * as websocketAction from '../websockets/websocket-action'
 
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
 
@@ -383,7 +383,7 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
       userId: TEST_USER_ID,
       clientSessionId: 'test-session-delete-function-integration',
       localAgentTemplates: mockLocalAgentTemplates,
-      onResponseChunk: (chunk: string | PrintModeObject) => {
+      onResponseChunk: (chunk: string | PrintModeEvent) => {
         if (typeof chunk !== 'string') {
           return
         }
@@ -479,7 +479,7 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
             stepPrompt: '',
           },
         },
-        onResponseChunk: (chunk: string | PrintModeObject) => {
+        onResponseChunk: (chunk: string | PrintModeEvent) => {
           if (typeof chunk !== 'string') {
             return
           }
diff --git a/backend/src/main-prompt.ts b/backend/src/main-prompt.ts
@@ -14,7 +14,7 @@ import { requestToolCall } from './websockets/websocket-action'
 import type { ClientToolCall } from './tools/constants'
 import type { ClientAction } from '@codebuff/common/actions'
 import type { CostMode } from '@codebuff/common/constants'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   SessionState,
   ToolResult,
@@ -25,7 +25,7 @@ import { AgentTemplate } from './templates/types'
 export interface MainPromptOptions {
   userId: string | undefined
   clientSessionId: string
-  onResponseChunk: (chunk: string | PrintModeObject) => void
+  onResponseChunk: (chunk: string | PrintModeEvent) => void
   localAgentTemplates: Record<string, AgentTemplate>
 }
 
diff --git a/backend/src/run-agent-step.ts b/backend/src/run-agent-step.ts
@@ -40,7 +40,7 @@ import { getRequestContext } from './websockets/request-context'
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { AgentResponseTrace } from '@codebuff/bigquery'
 import type { CodebuffMessage } from '@codebuff/common/types/message'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   AgentTemplateType,
   AgentState,
@@ -54,7 +54,7 @@ export interface AgentOptions {
   userInputId: string
   clientSessionId: string
   fingerprintId: string
-  onResponseChunk: (chunk: string | PrintModeObject) => void
+  onResponseChunk: (chunk: string | PrintModeEvent) => void
 
   agentType: AgentTemplateType
   fileContext: ProjectFileContext
@@ -516,7 +516,7 @@ export const loopAgentSteps = async (
 
     userId: string | undefined
     clientSessionId: string
-    onResponseChunk: (chunk: string | PrintModeObject) => void
+    onResponseChunk: (chunk: string | PrintModeEvent) => void
   },
 ) => {
   const agentTemplate = await getAgentTemplate(agentType, localAgentTemplates)
diff --git a/backend/src/run-programmatic-step.ts b/backend/src/run-programmatic-step.ts
@@ -12,7 +12,7 @@ import type {
   StepGenerator,
 } from '@codebuff/common/types/agent-template'
 import type { CodebuffToolCall } from './tools/constants'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   AgentState,
   AgentTemplateType,
@@ -63,7 +63,7 @@ export async function runProgrammaticStep(
     userInputId: string
     clientSessionId: string
     fingerprintId: string
-    onResponseChunk: (chunk: string | PrintModeObject) => void
+    onResponseChunk: (chunk: string | PrintModeEvent) => void
     agentType: AgentTemplateType
     fileContext: ProjectFileContext
     ws: WebSocket
diff --git a/backend/src/tools/handlers/tool/spawn-agents-async.ts b/backend/src/tools/handlers/tool/spawn-agents-async.ts
@@ -11,7 +11,7 @@ import type { CodebuffToolCall } from '../../constants'
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
 import type { SendSubagentChunk } from './spawn-agents'
 import type { CodebuffMessage } from '@codebuff/common/types/message'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   AgentState,
   AgentTemplateType,
@@ -129,7 +129,10 @@ export const handleSpawnAgentsAsync = ((params: {
     for (const { agent_type: agentTypeStr, prompt, params } of agents) {
       try {
         const agentType = agentTypeStr as AgentTemplateType
-        const agentTemplate = await getAgentTemplate(agentType, localAgentTemplates)
+        const agentTemplate = await getAgentTemplate(
+          agentType,
+          localAgentTemplates,
+        )
 
         if (!agentTemplate) {
           throw new Error(`Agent type ${agentTypeStr} not found.`)
@@ -205,7 +208,7 @@ export const handleSpawnAgentsAsync = ((params: {
               toolResults: [],
               userId,
               clientSessionId,
-              onResponseChunk: (chunk: string | PrintModeObject) => {
+              onResponseChunk: (chunk: string | PrintModeEvent) => {
                 if (typeof chunk !== 'string') {
                   return
                 }
diff --git a/backend/src/tools/handlers/tool/spawn-agents.ts b/backend/src/tools/handlers/tool/spawn-agents.ts
@@ -7,7 +7,7 @@ import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { CodebuffToolCall } from '../../constants'
 import type { CodebuffToolHandlerFunction } from '../handler-function-type'
 import type { CodebuffMessage } from '@codebuff/common/types/message'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   AgentState,
   AgentTemplateType,
@@ -113,7 +113,10 @@ export const handleSpawnAgents = ((params: {
     const results = await Promise.allSettled(
       agents.map(async ({ agent_type: agentTypeStr, prompt, params }) => {
         const agentType = agentTypeStr as AgentTemplateType
-        const agentTemplate = await getAgentTemplate(agentType, localAgentTemplates)
+        const agentTemplate = await getAgentTemplate(
+          agentType,
+          localAgentTemplates,
+        )
 
         if (!agentTemplate) {
           throw new Error(`Agent type ${agentTypeStr} not found.`)
@@ -191,7 +194,7 @@ export const handleSpawnAgents = ((params: {
           toolResults: [],
           userId,
           clientSessionId,
-          onResponseChunk: (chunk: string | PrintModeObject) => {
+          onResponseChunk: (chunk: string | PrintModeEvent) => {
             if (typeof chunk !== 'string') {
               return
             }
diff --git a/backend/src/tools/stream-parser.ts b/backend/src/tools/stream-parser.ts
@@ -11,7 +11,7 @@ import type { CodebuffToolCall } from './constants'
 import type { AgentTemplate } from '../templates/types'
 import type { ToolName } from '@codebuff/common/tools/constants'
 import type { CodebuffMessage } from '@codebuff/common/types/message'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   AgentState,
   Subgoal,
@@ -42,7 +42,7 @@ export async function processStreamWithTools<T extends string>(options: {
   messages: CodebuffMessage[]
   agentState: AgentState
   agentContext: Record<string, Subgoal>
-  onResponseChunk: (chunk: string | PrintModeObject) => void
+  onResponseChunk: (chunk: string | PrintModeEvent) => void
   fullResponse: string
 }) {
   const {
diff --git a/backend/src/tools/tool-executor.ts b/backend/src/tools/tool-executor.ts
@@ -14,7 +14,7 @@ import type { ClientToolCall, CodebuffToolCall } from './constants'
 import type { CodebuffToolHandlerFunction } from './handlers/handler-function-type'
 import type { AgentTemplate } from '../templates/types'
 import type { ToolName } from '@codebuff/common/tools/constants'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type { ToolResult } from '@codebuff/common/types/session-state'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 import type { WebSocket } from 'ws'
@@ -105,7 +105,7 @@ export interface ExecuteToolCallParams<T extends ToolName = ToolName> {
   clientSessionId: string
   userInputId: string
   fullResponse: string
-  onResponseChunk: (chunk: string | PrintModeObject) => void
+  onResponseChunk: (chunk: string | PrintModeEvent) => void
   state: Record<string, any>
   userId: string | undefined
   autoInsertEndStepParam?: boolean
diff --git a/common/src/actions.ts b/common/src/actions.ts
@@ -2,7 +2,7 @@ import { z } from 'zod'
 
 import { costModes } from './constants'
 import { GrantTypeValues } from './types/grant'
-import { printModeObjectSchema } from './types/print-mode'
+import { printModeEventSchema } from './types/print-mode'
 import {
   SessionStateSchema,
   toolCallSchema,
@@ -132,7 +132,7 @@ export const SERVER_ACTION_SCHEMA = z.discriminatedUnion('type', [
   z.object({
     type: z.literal('response-chunk'),
     userInputId: z.string(),
-    chunk: z.union([z.string(), printModeObjectSchema]),
+    chunk: z.union([z.string(), printModeEventSchema]),
   }),
   z.object({
     type: z.literal('subagent-response-chunk'),
diff --git a/common/src/types/print-mode.ts b/common/src/types/print-mode.ts
@@ -43,7 +43,7 @@ export const printModeFinishSchema = z.object({
 })
 export type PrintModeFinish = z.infer<typeof printModeFinishSchema>
 
-export const printModeObjectSchema = z.discriminatedUnion('type', [
+export const printModeEventSchema = z.discriminatedUnion('type', [
   printModeErrorSchema,
   printModeDownloadStatusSchema,
   printModeFinishSchema,
@@ -52,4 +52,4 @@ export const printModeObjectSchema = z.discriminatedUnion('type', [
   printModeToolResultSchema,
 ])
 
-export type PrintModeObject = z.infer<typeof printModeObjectSchema>
+export type PrintModeEvent = z.infer<typeof printModeEventSchema>
diff --git a/common/src/types/session-state.ts b/common/src/types/session-state.ts
@@ -87,7 +87,9 @@ export const AgentTemplateTypes = Object.fromEntries(
 ) as { [K in (typeof AgentTemplateTypeList)[number]]: UnderscoreToDash<K> }
 const agentTemplateTypeSchema = z.enum(AgentTemplateTypeList)
 // Allow dynamic agent types by extending the base enum with string
-export type AgentTemplateType = z.infer<typeof agentTemplateTypeSchema> | string
+export type AgentTemplateType =
+  | z.infer<typeof agentTemplateTypeSchema>
+  | (string & {})
 
 export const SessionStateSchema = z.object({
   fileContext: ProjectFileContextSchema,
diff --git a/npm-app/src/client.ts b/npm-app/src/client.ts
@@ -96,7 +96,7 @@ import type {
 } from '@codebuff/common/actions'
 import type { ApiKeyType } from '@codebuff/common/api-keys/constants'
 import type { CostMode } from '@codebuff/common/constants'
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 import type {
   SessionState,
   ToolResult,
@@ -1154,7 +1154,7 @@ export class Client {
   }
 
   private subscribeToResponse(
-    onChunk: (chunk: string | PrintModeObject) => void,
+    onChunk: (chunk: string | PrintModeEvent) => void,
     userInputId: string,
     onStreamStart: () => void,
     prompt: string,
diff --git a/npm-app/src/display/print-mode.ts b/npm-app/src/display/print-mode.ts
@@ -1,6 +1,6 @@
 import { originalConsoleError, originalConsoleLog } from './overrides'
 
-import type { PrintModeObject } from '@codebuff/common/types/print-mode'
+import type { PrintModeEvent } from '@codebuff/common/types/print-mode'
 
 let printModeEnabled: boolean = false
 export function setPrintMode(enabled: boolean) {
@@ -10,7 +10,7 @@ export function printModeIsEnabled(): boolean {
   return printModeEnabled ?? false
 }
 
-export function printModeLog(obj: PrintModeObject) {
+export function printModeLog(obj: PrintModeEvent) {
   if (!printModeEnabled) {
     return
   }
diff --git a/sdk/src/client.ts b/sdk/src/client.ts
@@ -1,30 +1,90 @@
+import { execFileSync } from 'child_process'
+
+import { processStream } from './process-stream'
+
 import type {
   CodebuffClientOptions,
-  Context,
+  ChatContext,
   ContinueChatOptions,
   NewChatOptions,
 } from './types'
 
+const CODEBUFF_BINARY = 'codebuff'
+
 export class CodebuffClient {
   private authToken: string
-  private abortController: AbortController
 
   public cwd: string
 
-  constructor({ apiKey, cwd, abortController }: CodebuffClientOptions) {
+  constructor({ apiKey, cwd }: 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"',
+      )
+    }
+
     this.authToken =
       apiKey.type === 'string'
         ? apiKey.value
         : process.env.CODEBUFF_API_KEY ?? ''
     this.cwd = cwd
-    this.abortController = abortController
   }
 
-  public newChat(options: NewChatOptions): Context {
-    return {}
+  public async runNewChat({
+    agent,
+    input: { prompt, params },
+    handleEvent,
+  }: NewChatOptions): Promise<ChatContext> {
+    const args = ['-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,
+      authToken: this.authToken,
+      handleEvent,
+    })
+
+    return {
+      agentId: agent,
+    }
   }
 
-  public continueChat(options: ContinueChatOptions): Context {
-    return {}
+  // WIP
+  private async continueChat({
+    agent,
+    input: { prompt, params },
+    context,
+    handleEvent,
+  }: ContinueChatOptions): Promise<ChatContext> {
+    agent = agent ?? context.agentId
+    const args = ['-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,
+      authToken: this.authToken,
+      handleEvent,
+    })
+
+    return {
+      agentId: agent,
+    }
   }
 }
diff --git a/sdk/src/process-stream.ts b/sdk/src/process-stream.ts
diff --git a/sdk/src/types.ts b/sdk/src/types.ts
