Don't enforce prompt response schema in client/sdk

Author: jahooma

npm-app/src/client.ts

@@ -1452,31 +1452,15 @@ export class Client {
     unsubscribeComplete = this.webSocket.subscribe(
       'prompt-response',
       async (action) => {
-        const parsedAction = PromptResponseSchema.safeParse(action)
-        if (!parsedAction.success) {
-          const message = [
-            'Received invalid prompt response from server:',
-            JSON.stringify(parsedAction.error.issues),
-            'If this issues persists, please contact support@codebuff.com',
-          ].join('\n')
-          console.error(message)
-          logger.error(
-            {
-              errorMessage: message,
-              action,
-              eventId: AnalyticsEvent.MALFORMED_PROMPT_RESPONSE,
-            },
-            'Malformed prompt response',
-          )
-          return
-        }
+        // Stop enforcing prompt response schema (e.g. PromptResponseSchema.parse(action))!
+        // It's a black box we will pass back to the server.
+
         if (action.promptId !== userInputId) return
-        const a = parsedAction.data
         this.responseComplete = true
 
         Spinner.get().stop()
 
-        this.sessionState = a.sessionState
+        this.sessionState = action.sessionState
         const toolResults: ToolResultPart[] = []
 
         stepsCount++
@@ -1543,8 +1527,8 @@ Go to https://www.codebuff.com/config for more information.`) +
         }
 
         // Print structured output as JSON if available
-        if (a.output?.type === 'structuredOutput') {
-          console.log('\n' + JSON.stringify(a.output.value, null, 2))
+        if (action.output?.type === 'structuredOutput') {
+          console.log('\n' + JSON.stringify(action.output.value, null, 2))
         }
 
         if (DiffManager.getChanges().length > 0) {
@@ -1586,7 +1570,7 @@ Go to https://www.codebuff.com/config for more information.`) +
         // Clear the onChunk callback when response is complete
         this.currentOnChunk = undefined
 
-        resolveResponse({ ...a, wasStoppedByUser: false })
+        resolveResponse({ ...action, wasStoppedByUser: false })
       },
     )
 

sdk/src/run.ts

@@ -32,7 +32,10 @@ import type {
   ToolResultPart,
 } from '../../common/src/types/messages/content-part'
 import type { PrintModeEvent } from '../../common/src/types/print-mode'
-import type { SessionState } from '../../common/src/types/session-state'
+import {
+  AgentOutputSchema,
+  type SessionState,
+} from '../../common/src/types/session-state'
 
 export type CodebuffClientOptions = {
   // Provide an API key or set the CODEBUFF_API_KEY environment variable.
@@ -360,27 +363,23 @@ async function handlePromptResponse({
       },
     })
   } else if (action.type === 'prompt-response') {
-    const parsedAction = PromptResponseSchema.safeParse(action)
-    if (!parsedAction.success) {
+    // Stop enforcing session state schema! It's a black box we will pass back to the server.
+    // Only check the output schema.
+    const parsedOutput = AgentOutputSchema.safeParse(action.output)
+    if (!parsedOutput.success) {
       const message = [
         'Received invalid prompt response from server:',
-        JSON.stringify(parsedAction.error.issues),
+        JSON.stringify(parsedOutput.error.issues),
         'If this issues persists, please contact support@codebuff.com',
       ].join('\n')
-      onError({
-        message: message,
-      })
+      onError({ message })
       resolve({
         sessionState: initialSessionState,
-        output: {
-          type: 'error',
-          message: message,
-        },
+        output: { type: 'error', message },
       })
       return
     }
-
-    const { sessionState, output } = parsedAction.data
+    const { sessionState, output } = action
     const state: RunState = {
       sessionState,
       output: output ?? {