remove zod schema for server action

Author: charleslien

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 }),
 )

common/src/actions.ts

@@ -1,23 +1,16 @@
 import { z } from 'zod/v4'
 
-import { GrantTypeValues } from './types/grant'
-import { mcpConfigSchema } from './types/mcp'
-import { toolMessageSchema } from './types/messages/codebuff-message'
-import { printModeEventSchema } from './types/print-mode'
-import {
-  AgentOutputSchema,
-  SessionStateSchema,
-  toolCallSchema,
-} from './types/session-state'
-
 import type { CostMode } from './old-constants'
+import type { GrantType } from './types/grant'
+import type { MCPConfig } from './types/mcp'
 import type { ToolMessage } from './types/messages/codebuff-message'
 import type {
   TextPart,
   ImagePart,
   ToolResultOutput,
 } from './types/messages/content-part'
-import type { SessionState } from './types/session-state'
+import type { PrintModeEvent } from './types/print-mode'
+import type { AgentOutput, SessionState, ToolCall } from './types/session-state'
 import type { ProjectFileContext } from './util/file'
 
 export const FileChangeSchema = z.object({
@@ -98,120 +91,125 @@ export type ClientAction<T extends ClientActionType = ClientActionType> = {
   >
 }[T]
 
-export const UsageReponseSchema = z.object({
-  type: z.literal('usage-response'),
-  usage: z.number(),
-  remainingBalance: z.number(),
-  balanceBreakdown: z
-    .record(
-      z.enum([GrantTypeValues[0], ...GrantTypeValues.slice(1)]),
-      z.number(),
-    )
-    .optional(),
-  next_quota_reset: z.coerce.date().nullable(),
-  autoTopupAdded: z.number().optional(),
-})
-export type UsageResponse = z.infer<typeof UsageReponseSchema>
-
-export const InitResponseSchema = z
-  .object({
-    type: z.literal('init-response'),
-    message: z.string().optional(),
-    agentNames: z.record(z.string(), z.string()).optional(),
-  })
-  .merge(
-    UsageReponseSchema.omit({
-      type: true,
-    }),
-  )
-export type InitResponse = z.infer<typeof InitResponseSchema>
-
-export const MessageCostResponseSchema = z.object({
-  type: z.literal('message-cost-response'),
-  promptId: z.string(),
-  credits: z.number(),
-  agentId: z.string().optional(),
-})
-export type MessageCostResponse = z.infer<typeof MessageCostResponseSchema>
-
-export const PromptResponseSchema = z.object({
-  type: z.literal('prompt-response'),
-  promptId: z.string(),
-  sessionState: SessionStateSchema,
-  toolCalls: z.array(toolCallSchema).optional(),
-  toolResults: z.array(toolMessageSchema).optional(),
-  output: AgentOutputSchema.optional(),
-})
-export type PromptResponse = z.infer<typeof PromptResponseSchema>
-
-export const SERVER_ACTION_SCHEMA = z.discriminatedUnion('type', [
-  z.object({
-    type: z.literal('response-chunk'),
-    userInputId: z.string(),
-    chunk: z.union([z.string(), printModeEventSchema]),
-  }),
-  z.object({
-    type: z.literal('subagent-response-chunk'),
-    userInputId: z.string(),
-    agentId: z.string(),
-    agentType: z.string(),
-    chunk: z.string(),
-    prompt: z.string().optional(),
-    forwardToPrompt: z.boolean().optional(),
-  }),
-  z.object({
-    type: z.literal('handlesteps-log-chunk'),
-    userInputId: z.string(),
-    agentId: z.string(),
-    level: z.enum(['debug', 'info', 'warn', 'error']),
-    data: z.any(),
-    message: z.string().optional(),
-  }),
-  PromptResponseSchema,
-  z.object({
-    type: z.literal('read-files'),
-    filePaths: z.array(z.string()),
-    requestId: z.string(),
-  }),
-  z.object({
-    type: z.literal('tool-call-request'),
-    userInputId: z.string(),
-    requestId: z.string(),
-    toolName: z.string(),
-    input: z.record(z.string(), z.any()),
-    timeout: z.number().optional(),
-    mcpConfig: mcpConfigSchema.optional(),
-  }),
-  InitResponseSchema,
-  UsageReponseSchema,
-  MessageCostResponseSchema,
-
-  z.object({
-    type: z.literal('action-error'),
-    message: z.string(),
-    error: z.string().optional(),
-    remainingBalance: z.number().optional(),
-  }),
-  z.object({
-    type: z.literal('prompt-error'),
-    userInputId: z.string(),
-    message: z.string(),
-    error: z.string().optional(),
-    remainingBalance: z.number().optional(),
-  }),
-  z.object({
-    // The server is imminently going to shutdown, and the client should reconnect
-    type: z.literal('request-reconnect'),
-  }),
-  z.object({
-    type: z.literal('request-mcp-tool-data'),
-    requestId: z.string(),
-    mcpConfig: mcpConfigSchema,
-    toolNames: z.string().array().optional(),
-  }),
-])
-
-type ServerActionAny = z.infer<typeof SERVER_ACTION_SCHEMA>
-export type ServerAction<
-  T extends ServerActionAny['type'] = ServerActionAny['type'],
-> = Extract<ServerActionAny, { type: T }>
+type ServerActionResponseChunk = {
+  type: 'response-chunk'
+  userInputId: string
+  chunk: string | PrintModeEvent
+}
+
+type ServerActionSubagentResponseChunk = {
+  type: 'subagent-response-chunk'
+  userInputId: string
+  agentId: string
+  agentType: string
+  chunk: string
+  prompt?: string
+  forwardToPrompt?: boolean
+}
+
+type ServerActionHandleStepsLogChunk = {
+  type: 'handlesteps-log-chunk'
+  userInputId: string
+  agentId: string
+  level: 'debug' | 'info' | 'warn' | 'error'
+  data: any
+  message?: string
+}
+
+export type PromptResponse = {
+  type: 'prompt-response'
+  promptId: string
+  sessionState: SessionState
+  toolCalls?: ToolCall[]
+  toolResults?: ToolMessage[]
+  output?: AgentOutput
+}
+
+type ServerActionReadFiles = {
+  type: 'read-files'
+  filePaths: string[]
+  requestId: string
+}
+
+type ServerActionToolCallRequest = {
+  type: 'tool-call-request'
+  userInputId: string
+  requestId: string
+  toolName: string
+  input: Record<string, any>
+  timeout?: number
+  mcpConfig?: MCPConfig
+}
+
+export type InitResponse = {
+  type: 'init-response'
+  message?: string
+  agentNames?: Record<string, string>
+} & Omit<UsageResponse, 'type'>
+
+export type UsageResponse = {
+  type: 'usage-response'
+  usage: number
+  remainingBalance: number
+  balanceBreakdown?: Record<GrantType, number>
+  next_quota_reset: Date | null
+  autoTopupAdded?: number
+}
+
+export type MessageCostResponse = {
+  type: 'message-cost-response'
+  promptId: string
+  credits: number
+  agentId?: string
+}
+
+type ServerActionActionError = {
+  type: 'action-error'
+  message: string
+  error?: string
+  remainingBalance?: number
+}
+
+type ServerActionPromptError = {
+  type: 'prompt-error'
+  userInputId: string
+  message: string
+  error?: string
+  remainingBalance?: number
+}
+
+type ServerActionRequestReconnect = {
+  // The server is imminently going to shutdown, and the client should reconnect
+  type: 'request-reconnect'
+}
+
+type ServerActionRequestMcpToolData = {
+  type: 'request-mcp-tool-data'
+  requestId: string
+  mcpConfig: MCPConfig
+  toolNames?: string[]
+}
+
+type ServerActionAny =
+  | ServerActionResponseChunk
+  | ServerActionSubagentResponseChunk
+  | ServerActionHandleStepsLogChunk
+  | PromptResponse
+  | ServerActionReadFiles
+  | ServerActionToolCallRequest
+  | InitResponse
+  | UsageResponse
+  | MessageCostResponse
+  | ServerActionActionError
+  | ServerActionPromptError
+  | ServerActionRequestReconnect
+  | ServerActionRequestMcpToolData
+type ServerActionType = ServerActionAny['type']
+export type ServerAction<T extends ServerActionType = ServerActionType> = {
+  [K in ServerActionType]: Extract<
+    ServerActionAny,
+    {
+      type: K
+    }
+  >
+}[T]