keep pending messages after cancel

Author: charleslien

cli/src/hooks/use-send-message.ts

@@ -1580,6 +1580,8 @@ export const useSendMessage = ({
           },
         })
 
+        previousRunStateRef.current = runState
+
         if (!runState.output || runState.output.type === 'error') {
           logger.warn(
             {
@@ -1625,8 +1627,6 @@ export const useSendMessage = ({
             }
           }),
         )
-
-        previousRunStateRef.current = runState
       } catch (error) {
         logger.error(
           { error: getErrorObject(error) },

packages/agent-runtime/src/util/messages.ts

@@ -68,6 +68,29 @@ export function buildUserMessageContent(
   return asUserMessage(textParts.join('\n\n'))
 }
 
+export function getCancelledAdditionalMessages(args: {
+  prompt: string | undefined
+  params: Record<string, any> | undefined
+  content?: Array<TextPart | ImagePart>
+  pendingAgentResponse: string
+}): Message[] {
+  const { prompt, params, content, pendingAgentResponse } = args
+
+  const messages: Message[] = [
+    {
+      role: 'user',
+      content: buildUserMessageContent(prompt, params, content),
+      tags: ['USER_PROMPT'],
+    },
+    {role: 'user', content: [
+      {type: 'text', text: `<previous_assistant_message>${pendingAgentResponse}</previous_assistant_message>`},
+      {type: 'text', text: asSystemMessage('Response cancelled by user.')},
+    ]},
+  ]
+
+  return messages
+}
+
 export function parseUserMessage(str: string): string | undefined {
   const match = str.match(/<user_message>(.*?)<\/user_message>/s)
   return match ? match[1] : undefined

sdk/src/run.ts

@@ -1,6 +1,7 @@
 import path from 'path'
 
 import { callMainPrompt } from '@codebuff/agent-runtime/main-prompt'
+import { getCancelledAdditionalMessages } from '@codebuff/agent-runtime/util/messages'
 import { MAX_AGENT_STEPS_DEFAULT } from '@codebuff/common/constants/agents'
 import { getMCPClient, listMCPTools } from '@codebuff/common/mcp/client'
 import { toOptionalFile } from '@codebuff/common/old-constants'
@@ -100,17 +101,6 @@ export type RunOptions = {
   signal?: AbortSignal
 }
 
-class AbortError extends Error {
-  name = 'AbortError'
-}
-
-function checkAborted(signal?: AbortSignal): AbortError | null {
-  if (!signal?.aborted) {
-    return null
-  }
-  return new AbortError('Run cancelled by user')
-}
-
 type RunReturnType = Awaited<ReturnType<typeof run>>
 export async function run({
   apiKey,
@@ -180,18 +170,6 @@ export async function run({
       logger,
     })
   }
-  {
-    const aborted = checkAborted(signal)
-    if (aborted) {
-      return {
-        sessionState,
-        output: {
-          type: 'error',
-          message: aborted.message,
-        },
-      }
-    }
-  }
 
   let resolve: (value: RunReturnType) => any = () => {}
   const promise = new Promise<RunReturnType>((res) => {
@@ -204,23 +182,50 @@ export async function run({
     }
   }
 
+  let pendingAgentResponse = ''
+  /** Calculates the current session state if cancelled.
+   *
+   * This includes the user'e message and pending assistant message.
+   */
+  function getCancelledSessionState(): SessionState {
+    const state = cloneDeep(sessionState)
+    state.mainAgentState.messageHistory.push(
+      ...getCancelledAdditionalMessages({
+        prompt,
+        params,
+        pendingAgentResponse,
+      }),
+    )
+    return state
+  }
+  function getCancelledRunState(): RunState {
+    return {
+      sessionState: getCancelledSessionState(),
+      output: {
+        type: 'error',
+        message: 'Run cancelled by user',
+      },
+    }
+  }
+
   const buffers: Record<string | 0, string> = { 0: '' }
 
   const onResponseChunk = async (
     action: ServerAction<'response-chunk'>,
   ): Promise<void> => {
-    const aborted = checkAborted(signal)
-    if (aborted) {
-      resolve({
-        sessionState,
-        output: {
-          type: 'error',
-          message: aborted.message,
-        },
-      })
+    if (signal?.aborted) {
       return
     }
     const { chunk } = action
+    addToPendingAssistantMessage: if (typeof chunk === 'string') {
+      pendingAgentResponse += chunk
+    } else if (
+      chunk.type === 'reasoning_delta' &&
+      chunk.ancestorRunIds.length === 0
+    ) {
+      pendingAgentResponse += chunk.text
+    }
+
     if (typeof chunk !== 'string') {
       if (chunk.type === 'reasoning_delta') {
         handleStreamChunk?.({
@@ -256,15 +261,7 @@ export async function run({
   const onSubagentResponseChunk = async (
     action: ServerAction<'subagent-response-chunk'>,
   ) => {
-    const aborted = checkAborted(signal)
-    if (aborted) {
-      resolve({
-        sessionState,
-        output: {
-          type: 'error',
-          message: aborted.message,
-        },
-      })
+    if (signal?.aborted) {
       return
     }
     const { agentId, agentType, chunk } = action
@@ -406,19 +403,6 @@ export async function run({
   const promptId = Math.random().toString(36).substring(2, 15)
 
   // Send input
-  {
-    const aborted = checkAborted(signal)
-    if (aborted) {
-      return {
-        sessionState,
-        output: {
-          type: 'error',
-          message: aborted.message,
-        },
-      }
-    }
-  }
-
   const userInfo = await getUserInfoFromApiKey({
     ...agentRuntimeImpl,
     apiKey,
@@ -429,6 +413,13 @@ export async function run({
   }
   const userId = userInfo.id
 
+  signal?.addEventListener('abort', () => {
+    resolve(getCancelledRunState())
+  })
+  if (signal?.aborted) {
+    return getCancelledRunState()
+  }
+
   callMainPrompt({
     ...agentRuntimeImpl,
     promptId,