use generic ServerAction and ClientAction instead of Extract

Author: charleslien

backend/src/loop-main-prompt.ts

@@ -13,7 +13,7 @@ const DEFAULT_MAX_ITERATIONS = 20
 
 export async function loopMainPrompt(
   ws: WebSocket,
-  action: Extract<ClientAction, { type: 'prompt' }>,
+  action: ClientAction<'prompt'>,
   options: MainPromptOptions & { maxIterations?: number },
 ): Promise<{
   sessionState: SessionState
@@ -32,7 +32,7 @@ export async function loopMainPrompt(
     toolCalls.length > 0 &&
     !toolCalls.some((tc) => tc.toolName === 'end_turn')
   ) {
-    const nextAction: Extract<ClientAction, { type: 'prompt' }> = {
+    const nextAction: ClientAction<'prompt'> = {
       ...action,
       sessionState,
       toolResults,

backend/src/main-prompt.ts

@@ -1,8 +1,6 @@
-import { uniq } from 'lodash'
-import type { WebSocket } from 'ws'
-
 import { renderToolResults } from '@codebuff/common/tools/utils'
 import { AgentTemplateTypes } from '@codebuff/common/types/session-state'
+import { uniq } from 'lodash'
 
 import { checkTerminalCommand } from './check-terminal-command'
 import { loopAgentSteps } from './run-agent-step'
@@ -21,6 +19,7 @@ import type {
   ToolResult,
   AgentTemplateType,
 } from '@codebuff/common/types/session-state'
+import type { WebSocket } from 'ws'
 
 export interface MainPromptOptions {
   userId: string | undefined
@@ -31,7 +30,7 @@ export interface MainPromptOptions {
 
 export const mainPrompt = async (
   ws: WebSocket,
-  action: Extract<ClientAction, { type: 'prompt' }>,
+  action: ClientAction<'prompt'>,
   options: MainPromptOptions,
 ): Promise<{
   sessionState: SessionState

backend/src/websockets/middleware.ts

@@ -53,7 +53,7 @@ export class WebSocketMiddleware {
 
   use<T extends ClientAction['type']>(
     callback: (
-      action: Extract<ClientAction, { type: T }>,
+      action: ClientAction<T>,
       clientSessionId: string,
       ws: WebSocket,
       userInfo: UserInfo | undefined,
@@ -100,14 +100,14 @@ export class WebSocketMiddleware {
 
   run<T extends ClientAction['type']>(
     baseAction: (
-      action: Extract<ClientAction, { type: T }>,
+      action: ClientAction<T>,
       clientSessionId: string,
       ws: WebSocket,
     ) => void,
     options: { silent?: boolean } = {},
   ) {
     return async (
-      action: Extract<ClientAction, { type: T }>,
+      action: ClientAction<T>,
       clientSessionId: string,
       ws: WebSocket,
     ) => {

backend/src/websockets/websocket-action.ts

@@ -132,7 +132,7 @@ export async function genUsageResponse(
  * @param ws - The WebSocket connection
  */
 const onPrompt = async (
-  action: Extract<ClientAction, { type: 'prompt' }>,
+  action: ClientAction<'prompt'>,
   clientSessionId: string,
   ws: WebSocket,
 ) => {
@@ -187,7 +187,7 @@ const onPrompt = async (
 
 export const callMainPrompt = async (
   ws: WebSocket,
-  action: Extract<ClientAction, { type: 'prompt' }>,
+  action: ClientAction<'prompt'>,
   options: {
     userId: string
     promptId: string
@@ -246,11 +246,7 @@ export const callMainPrompt = async (
  * @param ws - The WebSocket connection
  */
 const onInit = async (
-  {
-    fileContext,
-    fingerprintId,
-    authToken,
-  }: Extract<ClientAction, { type: 'init' }>,
+  { fileContext, fingerprintId, authToken }: ClientAction<'init'>,
   clientSessionId: string,
   ws: WebSocket,
 ) => {
@@ -284,7 +280,7 @@ const onInit = async (
 const onCancelUserInput = async ({
   authToken,
   promptId,
-}: Extract<ClientAction, { type: 'cancel-user-input' }>) => {
+}: ClientAction<'cancel-user-input'>) => {
   const userId = await getUserIdFromAuthToken(authToken)
   if (!userId) {
     logger.error({ authToken }, 'User id not found for authToken')
@@ -313,7 +309,7 @@ const callbacksByAction = {} as Record<
 export const subscribeToAction = <T extends ClientAction['type']>(
   type: T,
   callback: (
-    action: Extract<ClientAction, { type: T }>,
+    action: ClientAction<T>,
     clientSessionId: string,
     ws: WebSocket,
   ) => void,

common/src/websockets/websocket-client.ts

@@ -235,7 +235,7 @@ export class APIRealtimeClient {
 
   subscribe<T extends ServerAction['type']>(
     action: T,
-    callback: (action: Extract<ServerAction, { type: T }>) => void,
+    callback: (action: ServerAction<T>) => void,
   ) {
     const currSubscribers = this.subscribers.get(action) ?? []
     this.subscribers.set(action, [

npm-app/src/client.ts

@@ -1575,7 +1575,7 @@ Go to https://www.codebuff.com/config for more information.`) +
       this.setUsage(parsedAction.data)
     })
 
-    const initAction: Extract<ClientAction, { type: 'init' }> = {
+    const initAction: ClientAction<'init'> = {
       type: 'init',
       fingerprintId: await this.fingerprintId,
       authToken: this.user?.authToken,

sdk/src/client.ts

@@ -12,9 +12,9 @@ import {
 import { API_KEY_ENV_VAR } from '../../common/src/constants'
 import { getInitialSessionState } from '../../common/src/types/session-state'
 
+import type { AgentConfig } from './types/agent-config'
 import type { PrintModeEvent } from '../../common/src/types/print-mode'
 import type { SessionState } from '../../common/src/types/session-state'
-import type { AgentConfig } from './types/agent-config'
 
 type ClientToolName = 'write_file' | 'run_terminal_command'
 
@@ -27,7 +27,7 @@ export type CodebuffClientOptions = {
     Record<
       ClientToolName,
       (
-        args: Extract<ServerAction, { type: 'tool-call-request' }>['args'],
+        args: ServerAction<'tool-call-request'>['args'],
       ) => Promise<{ toolResultMessage: string }>
     > & {
       // Include read_files separately, since it has a different signature.
@@ -40,7 +40,7 @@ export type CodebuffClientOptions = {
 
 type RunState = {
   sessionState: SessionState
-  toolResults: Extract<ServerAction, { type: 'prompt-response' }>['toolResults']
+  toolResults: ServerAction<'prompt-response'>['toolResults']
 }
 
 export class CodebuffClient {
@@ -176,9 +176,7 @@ export class CodebuffClient {
     })
   }
 
-  private async handlePromptResponse(
-    action: Extract<ServerAction, { type: 'prompt-response' }>,
-  ) {
+  private async handlePromptResponse(action: ServerAction<'prompt-response'>) {
     const promiseActions =
       this.promptIdToResolveResponse[action?.promptId ?? '']
 
@@ -217,9 +215,7 @@ export class CodebuffClient {
     return getFiles(filePath, this.cwd)
   }
 
-  private async handleToolCall(
-    action: Extract<ServerAction, { type: 'tool-call-request' }>,
-  ) {
+  private async handleToolCall(action: ServerAction<'tool-call-request'>) {
     const toolName = action.toolName
     const args = action.args
     let result: string

sdk/src/websocket-client.ts

@@ -8,34 +8,23 @@ export type WebSocketHandlerOptions = {
   onWebsocketError?: (error: WebSocket.ErrorEvent) => void
   onWebsocketReconnect?: () => void
   onRequestReconnect?: () => Promise<void>
-  onResponseError?: (
-    error: Extract<ServerAction, { type: 'action-error' }>,
-  ) => Promise<void>
+  onResponseError?: (error: ServerAction<'action-error'>) => Promise<void>
   readFiles: (
     filePath: string[],
-  ) => Promise<Extract<ClientAction, { type: 'read-files-response' }>['files']>
+  ) => Promise<ClientAction<'read-files-response'>['files']>
   handleToolCall: (
-    action: Extract<ServerAction, { type: 'tool-call-request' }>,
-  ) => Promise<
-    Omit<
-      Extract<ClientAction, { type: 'tool-call-response' }>,
-      'type' | 'requestId'
-    >
-  >
+    action: ServerAction<'tool-call-request'>,
+  ) => Promise<Omit<ClientAction<'tool-call-response'>, 'type' | 'requestId'>>
   onCostResponse?: (
-    action: Extract<ServerAction, { type: 'message-cost-response' }>,
+    action: ServerAction<'message-cost-response'>,
   ) => Promise<void>
 
-  onResponseChunk?: (
-    action: Extract<ServerAction, { type: 'response-chunk' }>,
-  ) => Promise<void>
+  onResponseChunk?: (action: ServerAction<'response-chunk'>) => Promise<void>
   onSubagentResponseChunk?: (
-    action: Extract<ServerAction, { type: 'subagent-response-chunk' }>,
+    action: ServerAction<'subagent-response-chunk'>,
   ) => Promise<void>
 
-  onPromptResponse?: (
-    action: Extract<ServerAction, { type: 'prompt-response' }>,
-  ) => Promise<void>
+  onPromptResponse?: (action: ServerAction<'prompt-response'>) => Promise<void>
 
   apiKey: string
 }
@@ -160,7 +149,7 @@ export class WebSocketHandler {
 
   public sendInput(
     action: Omit<
-      Extract<ClientAction, { type: 'prompt' }>,
+      ClientAction<'prompt'>,
       keyof ReturnType<typeof this.getInputDefaultOptions>
     >,
   ) {